torchax 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

torchax/device_module.py

@@ -1,20 +1,33 @@
+import torch
+
+
 def _is_in_bad_fork():
   return False
 
+
 def manual_seed_all(seed):
   pass
 
+
 def device_count():
   return 1
 
+
 def get_rng_state():
   return []
 
+
 def set_rng_state(new_state, device):
   pass
 
+
 def is_available():
   return True
 
+
 def current_device():
-  return 0
+  return 0
+
+
+def get_amp_supported_dtype():
+  return [torch.float16, torch.bfloat16]

torchax/environment.py

@@ -1,2 +1 @@
 
-

torchax/export.py

@@ -4,14 +4,14 @@ import copy
 from typing import Any, Dict, Tuple
 import torch
 from torch.utils import _pytree as pytree
+import torchax
 from torchax import tensor
-from torchax.ops import ops_registry
+from torchax.ops import ops_registry, mappings
 from torchax import decompositions
 import jax
 import jax.export
 import sympy
 
-
 DEBUG = False
 
 
@@ -83,7 +83,8 @@ def exported_program_to_jax(exported_program, export_raw: bool = False):
   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)
+    exported_program = exported_program.run_decompositions(
+        decompositions.DECOMPOSITIONS)
   if DEBUG:
     print(exported_program.graph_module.code)
 
@@ -108,8 +109,8 @@ def exported_program_to_jax(exported_program, export_raw: bool = False):
 
   if export_raw:
     return names, states, func
-
-  states = pytree.tree_map_only(torch.Tensor, tensor.t2j, states)
+  env = torchax.default_env()
+  states = env.t2j_copy(states)
   return states, func
 
 
@@ -121,34 +122,35 @@ def extract_avals(exported):
   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)
+    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))
+    return jax.ShapeDtypeStruct(shape, mappings.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
+        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.
@@ -163,8 +165,10 @@ def extract_avals(exported):
         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}")
+      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)
@@ -182,7 +186,7 @@ def extract_avals(exported):
       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."
@@ -190,7 +194,8 @@ def extract_avals(exported):
       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)
+        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
@@ -203,8 +208,12 @@ def extract_avals(exported):
     # 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]
+    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

torchax/flax.py

@@ -0,0 +1,39 @@
+"""Flax interop."""
+
+import torch
+import torchax as tx
+import torchax.interop
+
+
+class FlaxNNModule(torch.nn.Module):
+
+  def __init__(self, env, flax_module, sample_args, sample_kwargs=None):
+    super().__init__()
+    prng = env.prng_key
+    sample_kwargs = sample_kwargs or {}
+    parameter_dict = tx.interop.call_jax(flax_module.init, prng, *sample_args,
+                                         **sample_kwargs)
+
+    self._params = self._encode_nested_dict(parameter_dict)
+
+    self._flax_module = flax_module
+
+  def _encode_nested_dict(self, nested_dict):
+    child_module = torch.nn.Module()
+    for k, v in nested_dict.items():
+      if isinstance(v, dict):
+        child_module.add_module(k, self._encode_nested_dict(v))
+      else:
+        child_module.register_parameter(k, torch.nn.Parameter(v))
+    return child_module
+
+  def _decode_nested_dict(self, child_module):
+    result = dict(child_module.named_parameters(recurse=False))
+    for k, v in child_module.named_children():
+      result[k] = self._decode_nested_dict(v)
+    return result
+
+  def forward(self, *args, **kwargs):
+    nested_dict_params = self._decode_nested_dict(self._params)
+    return tx.interop.call_jax(self._flax_module.apply, nested_dict_params,
+                               *args, **kwargs)