Apurva Mishra mav3ri3k

mav3ri3k / main.py

Created June 27, 2025 06:54

nnx.jit does not give error for impure functions. On bigger model, the difference between impure and pure function was way bigger.

	import jax
	import jax.numpy as jnp
	import flax.nnx as nnx
	import optax
	import time

	# change the train_step function definition and function call to see change
	# on line 53 and 78
	# On my machine
	# Pure function took 1.7 seconds

mav3ri3k / grpo_demo.py

Created January 29, 2025 12:45 — forked from willccbb/grpo_demo.py

GRPO Llama-1B

	# train_grpo.py
	import re
	import torch
	from datasets import load_dataset, Dataset
	from transformers import AutoTokenizer, AutoModelForCausalLM
	from peft import LoraConfig
	from trl import GRPOConfig, GRPOTrainer

	# Load and prep dataset

mav3ri3k / gist:88be5d6f1e8b9ddbe04d536c9a9d7e16

Created October 3, 2024 17:15

Final report for GSOC24 project: Sandboxed and Deterministic Proc Macro using Wasm

	# Sandboxed and Deterministic Proc Macro using Wasm

	This branch: 'gsoc24' is the final snapshot for state of code during the end of gsoc project: [Sandboxed and Deterministic Proc Macro using Wasm](https://summerofcode.withgoogle.com/programs/2024/projects/kXG0mZoj)
	The initial goal for the project was to:
	Add experimental support to rustc for building and running procedural macros as WebAssembly. Procedural Macro crates can opt in for being compiled to WebAssembly. This wasm-proc-macro will be a wasm blob sandboxed using WASM. It will interact with the compiler only through a stream of token and no ability to interact with the outside world.

	The project was inspired by: [Build-time execution sandboxing](https://github.com/rust-lang/compiler-team/issues/475) and [Pre-RFC: Sandboxed, deterministic, reproducible, efficient Wasm compilation of proc macros](https://internals.rust-lang.org/t/pre-rfc-sandboxed-deterministic-reproducible-efficient-wasm-compilation-of-proc-macros/19359)

	I started with the project