torchax 0.0.4__py3-none-any.whl

torchax/distributed.py

@@ -0,0 +1,246 @@
+"""`torch.distributed` backend implemented with JAX collective ops.
+
+EXPERIMENTAL: This module is still highly experimental, and it may be removed
+before any stable release.
+
+Note: JAX collective ops require that axis names be defined in `pmap` or
+`shmap`. The distributed backend only supports one axis, named `torch_dist`.
+This name is defined by our mirror implementation of `spawn`.
+"""
+
+import datetime
+import functools
+import logging
+import os
+from typing import List, Optional, Union
+
+import jax
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.distributed._functional_collectives
+from torch._C._distributed_c10d import ProcessGroup  # type: ignore
+import torch.distributed
+import torchax
+from jax.sharding import NamedSharding
+from jax.sharding import Mesh, PartitionSpec as P
+from jax.experimental import mesh_utils
+import torch.utils._pytree as torch_pytree
+from torchax import interop
+
+
+class ProcessGroupJax(ProcessGroup):
+  """Distributed backend implemented with JAX."""
+
+  def __init__(self, prefix_store, rank, size, timeout):
+    super().__init__(rank, size)
+    self._group_name = None
+
+  def getBackendName(self):
+    return "jax"
+
+  # TODO(wcromar): why doesn't default group name setter work?
+  # https://github.com/pytorch/pytorch/blob/7b1988f9222f3dec5cc2012afce84218199748ae/torch/csrc/distributed/c10d/ProcessGroup.cpp#L148-L152
+  def _set_group_name(self, name: str) -> None:
+    self._group_name = name
+
+  @property
+  def group_name(self):
+    assert self._group_name
+    return self._group_name
+
+  @staticmethod
+  def _work(
+    tensors: Union[torch.Tensor, List[torch.Tensor], List[List[torch.Tensor]]],
+  ) -> dist.Work:
+    fut = torch.futures.Future()
+    fut.set_result(tensors)
+    return torch._C._distributed_c10d._create_work_from_future(fut)
+
+  def _allgather_base(
+    self,
+    output: torch.Tensor,
+    input: torch.Tensor,
+    opts=...,
+  ) -> dist.Work:
+    assert isinstance(input, torchax.tensor.Tensor)
+    assert isinstance(output, torchax.tensor.Tensor)
+    torch.distributed._functional_collectives.all_gather_tensor_inplace(
+      output, input, group=self
+    )
+    return self._work(output)
+
+  def allreduce(
+    self,
+    tensors: List[torch.Tensor],
+    opts: dist.AllreduceOptions = ...,
+  ) -> dist.Work:
+    assert len(tensors) == 1
+    assert isinstance(tensors[0], torchax.tensor.Tensor)
+    torch.distributed._functional_collectives.all_reduce_inplace(
+      tensors[0],
+      torch.distributed._functional_collectives.REDUCE_OP_TO_STR[
+        opts.reduceOp.op
+      ],
+      self,
+    )
+
+    return self._work(tensors)
+
+  def broadcast(
+    self,
+    tensors: List[torch.Tensor],
+    opts: dist.BroadcastOptions = ...,
+  ) -> dist.Work:
+    assert len(tensors) == 1
+    assert isinstance(tensors[0], torchax.tensor.Tensor)
+    tensors[0].copy_(
+      torch.distributed._functional_collectives.broadcast(
+        tensors[0], opts.rootRank, group=self
+      )
+    )
+
+    return self._work(tensors)
+
+
+dist.Backend.register_backend("jax", ProcessGroupJax)
+
+
+def jax_rendezvous_handler(
+  url: str, timeout: datetime.timedelta = ..., **kwargs
+):
+  """Initialize distributed store with JAX process IDs.
+
+  Requires `$MASTER_ADDR` and `$MASTER_PORT`.
+  """
+  # TODO(wcromar): jax.distributed.initialize(...) for multiprocess on GPU
+  # TODO(wcromar): Can we use the XLA coordinator as a Store? This isn't part
+  # of their public Python API
+  master_ip = os.environ["MASTER_ADDR"]
+  master_port = int(os.environ["MASTER_PORT"])
+  # TODO(wcromar): Use `torchrun`'s store if available
+  store = dist.TCPStore(
+    master_ip,
+    master_port,
+    jax.process_count(),
+    is_master=jax.process_index() == 0,
+  )
+
+  yield (store, jax.process_index(), jax.process_count())
+
+
+dist.register_rendezvous_handler("jax", jax_rendezvous_handler)
+
+
+def spawn(f, args=(), env: Optional[torchax.tensor.Environment] = None):
+  """Wrap `f` in a JAX `pmap` with the axis name `torch_dist` defined.
+  `f` is expected to take the replica index as a positional argument, similar
+  to `torch.multiprocessing.spawn`.
+  Note: `spawn` does not actually create parallel processes.
+  """
+  env = env or torchax.default_env()
+
+  def jax_wrapper(index, jax_args):
+    index, args = env.j2t_iso([index, jax_args])
+    torch_outputs = f(index, *args)
+    return env.t2j_iso(torch_outputs)
+
+  jax_outputs = jax.pmap(jax_wrapper, axis_name="torch_dist")(
+    np.arange(jax.device_count()), env.t2j_iso(args)
+  )
+  return env.j2t_iso(jax_outputs)
+
+
+class DistributedDataParallel(torch.nn.Module):
+  """Re-implementation of DistributedDataParallel using JAX SPMD.
+
+  Splits inputs along batch dimension (assumed to be 0) across all devices in
+  JAX runtime, including remote devices. Each process should load a distinct
+  shard of the input data using e.g. DistributedSampler. Each process' shard
+  is then further split among the addressable devices (e.g. local TPU chips)
+  by `shard_input`.
+
+  Note: since parameters are replicated across addressable devices, inputs
+  must also be SPMD sharded using `shard_input` or `replicate_input`.
+
+  Example usage:
+
+  ```
+  jax_model = torchax.distributed.DistributedDataParallel(create_model())
+  for data, dataloader:
+    jax_data = jax_model.shard_input(data)
+    jax_output = jax_model(jax_data)
+  ```
+  """
+  def __init__(
+    self,
+    module: torch.nn.Module,
+    env: Optional[torchax.tensor.Environment] = None,
+    **kwargs,
+  ):
+    if kwargs:
+      logging.warning(f"Unsupported kwargs {kwargs}")
+
+    super().__init__()
+    self._env = env or torchax.default_env()
+    self._mesh = Mesh(
+      mesh_utils.create_device_mesh((jax.device_count(),)),
+      axis_names=("batch",),
+    )
+    replicated_state = torch_pytree.tree_map_only(
+      torch.Tensor,
+      lambda t: self._env.j2t_iso(
+        jax.device_put(
+          self._env.to_xla(t)._elem, NamedSharding(self._mesh, P())
+        )
+      ),
+      module.state_dict(),
+    )
+    # TODO: broadcast
+    module.load_state_dict(replicated_state, assign=True)
+    self._module = module
+
+  def shard_input(self, inp):
+    per_process_batch_size = inp.shape[0]  # assumes batch dim is 0
+    per_replica_batch_size = per_process_batch_size // jax.local_device_count()
+    per_replica_batches = torch.chunk(inp, jax.local_device_count())
+    global_batch_size = per_replica_batch_size * jax.device_count()
+    global_batch_shape = (global_batch_size,) + inp.shape[1:]
+
+    sharding = NamedSharding(self._mesh, P("batch"))
+    return self._env.j2t_iso(jax.make_array_from_single_device_arrays(
+      global_batch_shape,
+      NamedSharding(self._mesh, P("batch")),
+      arrays=[
+        jax.device_put(self._env.to_xla(batch)._elem, device)
+        for batch, device in zip(
+          per_replica_batches, sharding.addressable_devices
+        )
+      ],
+    ))
+
+  def replicate_input(self, inp):
+    return self._env.j2t_iso(
+      jax.device_put(inp._elem, NamedSharding(self._mesh, P()))
+    )
+
+  def jit_step(self, func):
+    @functools.partial(interop.jax_jit,
+                       kwargs_for_jax_jit={'donate_argnums': 0})
+    def _jit_fn(states, args):
+      self.load_state_dict(states)
+      outputs = func(*args)
+      return self.state_dict(), outputs
+
+    @functools.wraps(func)
+    def inner(*args):
+      jax_states = self.state_dict()
+      new_states, outputs = _jit_fn(jax_states, args)
+      self.load_state_dict(new_states)
+      return outputs
+
+    return inner
+
+  def forward(self, *args):
+    with self._env:
+      return self._module(*args)

torchax/environment.py

@@ -0,0 +1,2 @@
+
+

torchax/export.py

@@ -0,0 +1,236 @@
+# pylint: disable
+"""Utilities for exporting a torch program to jax/stablehlo."""
+import copy
+from typing import Any, Dict, Tuple
+import torch
+from torch.utils import _pytree as pytree
+from torchax import tensor
+from torchax.ops import ops_registry
+from torchax import decompositions
+import jax
+import jax.export
+import sympy
+
+
+DEBUG = False
+
+
+class JaxInterpreter(torch.fx.Interpreter):
+  """Experimental."""
+
+  def __init__(self, graph_module):
+    super().__init__(graph_module)
+    import torchax.ops.jaten
+    import torchax.ops.jtorch
+
+  def call_function(self, target, args: Tuple, kwargs: Dict) -> Any:
+    if not isinstance(target,
+                      (torch._ops.OpOverloadPacket, torch._ops.OpOverload)):
+      return super().call_function(target, args, kwargs)
+
+    if DEBUG:
+      print('Running ', target.name(), '--------')
+
+    op = ops_registry.all_aten_ops.get(target)
+    if op is None:
+      op = ops_registry.all_aten_ops.get(target.overloadpacket)
+    assert op is not None, target
+    assert op.is_jax_function, op
+    if op is None:
+      op = ops_registry.all_aten_ops.get(target.overloadpacket)
+    if op is None:
+      print(target.name(), target.tags)
+      raise RuntimeError('No lowering found for', target.name())
+    return op.func(*args, **kwargs)
+
+  def run_node(self, n) -> Any:
+    res = super().run_node(n)
+    if DEBUG:
+      if n.op == 'call_function':
+        if hasattr(res, 'shape'):
+          print('Meta:', n.meta.get('val').shape, 'REAL: ', res.shape)
+    return res
+
+
+from torch._decomp import get_decompositions
+import torch._refs
+
+_extra_decomp = get_decompositions([torch.ops.aten.unfold])
+
+
+def _extract_states_from_exported_program(exported_model):
+  # NOTE call convention: (parameters, buffers, user_inputs)
+  param_and_buffer_keys = exported_model.graph_signature.parameters + exported_model.graph_signature.buffers
+  state_dict = copy.copy(exported_model.state_dict)
+  if (constants := getattr(exported_model, 'constants', None)) is not None:
+    state_dict.update(constants)
+  param_buffer_values = list(state_dict[key] for key in param_and_buffer_keys)
+
+  if hasattr(exported_model.graph_signature, "lifted_tensor_constants"):
+    for name in exported_model.graph_signature.lifted_tensor_constants:
+      param_buffer_values.append(exported_model.tensor_constants[name])
+
+  return param_and_buffer_keys, param_buffer_values
+
+
+def exported_program_to_jax(exported_program, export_raw: bool = False):
+  """returns a pytree of jax arrays(state), and
+
+  a callable(func) that is jax function.
+
+  func(state, input) would be how you call it.
+  """
+  if torch.__version__ >= '2.2':
+    # torch version 2.1 didn't expose this yet
+    exported_program = exported_program.run_decompositions()
+    exported_program = exported_program.run_decompositions(decompositions.EXTRA_DECOMP)
+  if DEBUG:
+    print(exported_program.graph_module.code)
+
+  names, states = _extract_states_from_exported_program(exported_program)
+
+  def _extract_args(args, kwargs):
+    flat_args, received_spec = pytree.tree_flatten(
+        (args, kwargs))  # type: ignore[possibly-undefined]
+    return flat_args
+
+  num_mutations = len(exported_program.graph_signature.buffers_to_mutate)
+
+  def func(states, inputs):
+    args = _extract_args(inputs, {})
+    res = JaxInterpreter(exported_program.graph_module).run(
+        *states,
+        *args,
+        enable_io_processing=False,
+    )
+    res = res[num_mutations:]
+    return res
+
+  if export_raw:
+    return names, states, func
+
+  states = pytree.tree_map_only(torch.Tensor, tensor.t2j, states)
+  return states, func
+
+
+def extract_avals(exported):
+  """Return JAX Abstract Value shapes for all input parameters of the exported
+  program. This supports dynamic batch dimensions, including with constraints.
+  """
+
+  def _to_aval(arg_meta, symbolic_shapes):
+    """Convet from torch type to jax abstract value for export tracing
+    """
+    def _get_dim(d):
+      if isinstance(d, torch.SymInt):
+        return symbolic_shapes[str(d)]
+      return d
+
+    val = arg_meta['val']
+    is_scalar = isinstance(val, float) or isinstance(val, int) or isinstance(val, bool)
+    if is_scalar:
+      return jax.ShapeDtypeStruct([], type(arg_meta['val']))
+
+    tensor_meta = arg_meta['tensor_meta']
+    shape = [_get_dim(d) for d in tensor_meta.shape]
+    return jax.ShapeDtypeStruct(shape, tensor.t2j_dtype(tensor_meta.dtype))
+
+  def _get_inputs(exported):
+    """Return placeholders with input metadata"""
+    placeholders = [p for p in exported.graph.nodes if p.op == "placeholder"]
+    input_placeholders = [
+      p
+      for p, s in zip(placeholders, exported.graph_signature.input_specs)
+      if s.kind == torch.export.graph_signature.InputKind.USER_INPUT
+    ]
+    return input_placeholders
+
+  def _build_symbolic_shapes(range_constraints):
+    """Convert torch SymInt to JAX symbolic_shape and stores in a map using the
+    string name of the torch symbolic int.
+ 
+    TODO: There is probably a better way of storing a key for a symbolic int.
+    This value needs to be looked up again in `_to_aval` to figure out which
+    JAX symbolic to map to for a given torch tensor.
+    """
+    if len(range_constraints) == 0:
+      return None
+
+    def _build_symbolic_constraints(symbol_name, torch_constraint):
+      """Convert torch SymInt constraints to string for JAX symbolic_shape
+      Using sympy may be overkill here, currently PyTorch only uses ValueRanges
+      which allow specifying the min and the max of a value, for example:
+        torch.export.Dim("a", min=5, max=10)
+         ==> ("a >= 5", "a <= 10",)
+      """
+      if not isinstance(torch_constraint, torch.utils._sympy.value_ranges.ValueRanges) or torch_constraint.is_bool:
+        raise TypeError(f"No symbolic constraint handler for: {torch_constraint}")
+
+      constraints = []
+      symbol = sympy.Symbol(symbol_name)
+      if torch_constraint.lower != 2:
+        constraints.append(symbol >= torch_constraint.lower)
+      from sympy.core.singleton import S
+      if not torch_constraint.upper.is_infinite and torch_constraint.upper is not S.IntInfinity:
+        constraints.append(symbol <= torch_constraint.upper)
+
+      return tuple(sympy.pretty(c, use_unicode=False) for c in constraints)
+
+    def _build_symbolic_shape(sym, constraint, free_symbols):
+      """Returns a JAX symbolic shape for a given symbol and constraint
+
+      There are two possible sympy `sym` inputs:
+        1. Symbol - (s0) These can have custom constraints.
+        2. Expr - (s0*2) These apply the expr to s0's constraints, cannot override.
+      
+        Currently support is limited to operations with a symbol and and int,
+        in `torch/export/dynamic_shapes.py`:
+        "Only increasing linear operations with integer coefficients are supported."
+      """
+      symbol_name = str(sym)
+      constraints = _build_symbolic_constraints(symbol_name, constraint)
+      if sym.is_symbol:
+        symbolic_shape = jax.export.symbolic_shape(symbol_name, constraints=constraints)
+      else:
+        assert len(sym.free_symbols) > 0
+        scope = free_symbols[str(list(sym.free_symbols)[0])].scope
+        symbolic_shape = jax.export.symbolic_shape(symbol_name, scope=scope)
+      assert len(symbolic_shape) == 1
+      return symbolic_shape[0]
+
+    # Populate symbol variables before expressions, exprs need to use the same
+    # Symbolic scope as the variable they operate on. Expressions can only be
+    # integer compuations on symbol variables, so each symbol variable is OK to
+    # have its own scope.
+    symbolic_shapes = {}
+    symbol_variables = [(s,v) for s,v in range_constraints.items() if s.is_symbol]
+    symbol_exprs = [(s,v) for s,v in range_constraints.items() if not s.is_symbol]
+    for sym, constraint in symbol_variables + symbol_exprs:
+      symbolic_shape = _build_symbolic_shape(sym, constraint, symbolic_shapes)
+      symbolic_shapes[str(sym)] = symbolic_shape
+    return symbolic_shapes
+
+  symbolic_shapes = _build_symbolic_shapes(exported.range_constraints)
+  args = _get_inputs(exported)
+
+  if DEBUG:
+    print('Inputs to aval:', args, '--------')
+    print('Symbolic shapes:', symbolic_shapes)
+    for arg in args:
+      print('Meta2Aval', arg.meta, '--> ', _to_aval(arg.meta, symbolic_shapes))
+
+  return [_to_aval(arg.meta, symbolic_shapes) for arg in args]
+
+
+def exported_program_to_stablehlo(exported_program):
+  """Replacement for torch_xla.stablehlo.exported_program_to_stablehlo
+
+  Convert a program exported via torch.export to StableHLO.
+
+  This supports dynamic dimension sizes and generates explicit checks for
+  dynamo guards in the IR using shape_assertion custom_call ops.
+  """
+  weights, func = exported_program_to_jax(exported_program)
+  jax_avals = extract_avals(exported_program)
+  jax_export = jax.export.export(jax.jit(func))(weights, (jax_avals,))
+  return weights, jax_export

torchax/interop.py

@@ -0,0 +1,209 @@
+import copy
+import functools
+import torch
+from torch.nn.utils import stateless as torch_stateless
+import jax
+import jax.numpy as jnp
+from jax import tree_util as pytree
+from jax.experimental.shard_map import shard_map
+from torchax import tensor
+import torchax
+
+from torchax.types import JaxValue, TorchValue, JaxCallable, TorchCallable
+
+
+def extract_all_buffers(m: torch.nn.Module):
+    buffers = {}
+    params = {}
+    def extract_one(module, prefix):
+        for k in dir(module):
+            try:
+                v = getattr(module, k)
+            except:
+                continue
+            qual_name = prefix + k
+            if isinstance(v, torch.nn.parameter.Parameter) and v.requires_grad:
+                params[qual_name] = v
+            elif isinstance(v, torch.Tensor):
+                buffers[qual_name] = v
+        for name, child in module.named_children():
+            extract_one(child, prefix + name + '.')
+    extract_one(m, '')
+    return params, buffers
+
+
+def set_all_buffers(m, params, buffers):
+    def set_one(module, prefix):
+        for k in dir(module):
+            qual_name = prefix + k
+            if (potential_v := buffers.get(qual_name)) is not None:
+                setattr(module, k, potential_v)
+            elif (potential_v := params.get(qual_name)) is not None:
+                print(k, potential_v)
+                setattr(module, k, torch.nn.Parameter(potential_v))
+        for name, child in module.named_children():
+            set_one(child, prefix + name + '.')
+
+    set_one(m, '')
+
+
+class JittableModule(torch.nn.Module):
+
+    def __init__(self, m: torch.nn.Module, extra_jit_args={}):
+        super().__init__()
+        self.params, self.buffers = extract_all_buffers(m)
+        self._model = m
+        self._jitted = {}
+
+        self._extra_jit_args = extra_jit_args
+
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+
+    def functional_call(
+            self, method_name, params, buffers, *args, **kwargs):
+        kwargs = kwargs or {}
+        params_copy = copy.copy(params)
+        params_copy.update(buffers)
+        with torch_stateless._reparametrize_module(self._model, params_copy):
+            res = getattr(self._model, method_name)(*args, **kwargs)
+        return res
+
+
+    def forward(self, *args, **kwargs):
+        if 'forward' not in self._jitted:
+            jitted = jax_jit(
+                functools.partial(self.functional_call, 'forward'),
+                kwargs_for_jax_jit=self._extra_jit_args,
+            )
+            def jitted_forward(*args, **kwargs):
+                return jitted(self.params, self.buffers, *args, **kwargs)
+            self._jitted['forward'] = jitted_forward
+        return self._jitted['forward'](*args, **kwargs)
+
+    def __getattr__(self, key):
+        if key == '_model':
+            return super().__getattr__(key)
+        if key in self._jitted:
+            return self._jitted[key]
+        return getattr(self._model, key)
+
+    def make_jitted(self, key):
+        jitted = jax_jit(
+            functools.partial(self.functional_call, key), 
+            kwargs_for_jax_jit=self._extra_jit_args)
+        def call(*args, **kwargs):
+            return jitted(self.params, self.buffers, *args, **kwargs)
+        self._jitted[key] = call
+
+
+
+
+
+class CompileMixin:
+
+    def functional_call(
+            self, method, params, buffers, *args, **kwargs):
+        kwargs = kwargs or {}
+        params_copy = copy.copy(params)
+        params_copy.update(buffers)
+        with torch_stateless._reparametrize_module(self, params_copy):
+            res = method(*args, **kwargs)
+        return res
+
+    def jit(self, method):
+        jitted = jax_jit(functools.partial(self.functional_call, method_name))
+        def call(*args, **kwargs):
+            return jitted(self.named_paramters(), self.named_buffers(), *args, **kwargs)
+        return call
+
+
+def compile_nn_module(m: torch.nn.Module, methods=None):
+    if methods is None:
+        methods = ['forward']
+
+    new_parent = type(
+        m.__class__.__name__ + '_with_CompileMixin',
+        (CompileMixin, m.__class__),
+    )
+    m.__class__ = NewParent
+
+
+def _torch_view(t: JaxValue) -> TorchValue:
+    # t is an object from jax land
+    # view it as-if it's a torch land object
+    if isinstance(t, jax.Array):
+        # TODO
+        return tensor.Tensor(t, torchax.default_env())
+    if isinstance(t, type(jnp.int32)):
+        return tensor.t2j_type(t)
+    if callable(t):  # t is a JaxCallable
+        return functools.partial(call_jax, t)
+    # regular types are not changed
+    return t
+
+torch_view = functools.partial(pytree.tree_map, _torch_view)
+
+
+def _jax_view(t: TorchValue) -> JaxValue:
+    # t is an object from torch land
+    # view it as-if it's a jax land object
+    if isinstance(t, torch.Tensor):
+        assert isinstance(t, tensor.Tensor), type(t)
+        return t.jax()
+    if isinstance(t, type(torch.int32)):
+        return tensor.t2j_dtype(t)
+
+    # torch.nn.Module needs special handling
+    if not isinstance(t, torch.nn.Module) and callable(t):  # t is a TorchCallable
+        return functools.partial(call_torch, t)
+    # regular types are not changed
+    return t
+
+jax_view = functools.partial(pytree.tree_map, _jax_view)
+
+
+def call_jax(jax_func: JaxCallable, 
+             *args: TorchValue, 
+             **kwargs: TorchValue) -> TorchValue:
+    args, kwargs = jax_view((args, kwargs))
+    res: JaxValue = jax_func(*args, **kwargs)
+    return torch_view(res)
+
+
+def call_torch(torch_func: TorchCallable, *args: JaxValue, **kwargs: JaxValue) -> JaxValue:
+    args, kwargs = torch_view((args, kwargs))
+    with torchax.default_env():
+        res: TorchValue = torch_func(*args, **kwargs)
+    return jax_view(res)
+
+
+fori_loop = torch_view(jax.lax.fori_loop)
+
+
+def wrap_jax_jit(torch_function, jax_jit_func=jax.jit, kwargs_for_jax=None):
+    kwargs_for_jax = kwargs_for_jax or {}
+    jax_func = jax_view(torch_function)
+    jitted = jax_jit_func(jax_func, **kwargs_for_jax)
+    return torch_view(jitted)
+
+
+def jax_jit(torch_function, kwargs_for_jax_jit=None):
+    return wrap_jax_jit(torch_function, jax_jit_func=jax.jit,
+                        kwargs_for_jax=kwargs_for_jax_jit)
+
+
+def jax_shard_map(torch_function, kwargs_for_jax_shard_map=None):
+    return wrap_jax_jit(torch_function, jax_jit_func=shard_map,
+                        kwargs_for_jax=kwargs_for_jax_shard_map)
+
+
+def jax_value_and_grad(torch_function, kwargs_for_value_and_grad=None):
+    return wrap_jax_jit(torch_function, jax_jit_func=jax.value_and_grad,
+                        kwargs_for_jax=kwargs_for_value_and_grad)
+
+def gradient_checkpoint(torch_function, kwargs=None):
+    return wrap_jax_jit(torch_function, jax_jit_func=jax.checkpoint,
+                        kwargs_for_jax=kwargs)

torchax/ops/__init__.py

@@ -0,0 +1,10 @@
+def all_aten_jax_ops():
+  # to load the ops
+  import torchax.ops.jaten  # type: ignore
+  import torchax.ops.ops_registry  # type: ignore
+
+  return {
+      key: val.func
+      for key, val in torchax.ops.ops_registry.all_aten_ops.items()
+      if val.is_jax_function
+  }