torchax 0.0.10.dev20251114__py3-none-any.whl → 0.0.11.dev202612__py3-none-any.whl

torchax/__init__.py

@@ -13,16 +13,21 @@
 # limitations under the License.
 
 import contextlib
-from typing import List, Dict, Any, Optional
 import dataclasses
-import jax
 import os
+from contextlib import contextmanager
+from typing import Any
+
+import jax
 import torch
 from torch.utils import _pytree as pytree
+
+import torchax.device_module
 from torchax import tensor
-from contextlib import contextmanager
 
-__version__ = "0.0.10.dev20251114"
+from .checkpoint import load_checkpoint, save_checkpoint
+
+__version__ = "0.0.11.dev202612"
 VERSION = __version__
 
 # the "fast path" uses some sparse tensor thingies that currently we
@@ -31,123 +36,114 @@ torch.backends.mha.set_fastpath_enabled(False)
 
 
 __all__ = [
-    "default_env",
-    "extract_jax",
-    "enable_globally",
-    "save_checkpoint",
-    "load_checkpoint",
+  "default_env",
+  "extract_jax",
+  "enable_globally",
+  "save_checkpoint",
+  "load_checkpoint",
 ]
 
-from .checkpoint import save_checkpoint, load_checkpoint
 
 os.environ.setdefault("ENABLE_RUNTIME_UPTIME_TELEMETRY", "1")
 
 # torchax:oss-begin
 if getattr(jax.config, "jax_pjrt_client_create_options", None):
-    jax.config.update(
-        "jax_pjrt_client_create_options",
-        f"ml_framework_name:PyTorch/XLA2;ml_framework_version:{'v0.0.1'}",
-    )
+  jax.config.update(
+    "jax_pjrt_client_create_options",
+    f"ml_framework_name:PyTorch/XLA2;ml_framework_version:{'v0.0.1'}",
+  )
 # torchax:oss-end
 
 env = None
 
 
 def default_env():
-    global env
+  global env
 
-    if env is None:
-        env = tensor.Environment()
-    return env
+  if env is None:
+    env = tensor.Environment()
+  return env
 
 
 def extract_jax(mod: torch.nn.Module, env=None):
-    """Returns a pytree of jax.ndarray and a jax callable."""
-    if env is None:
-        env = default_env()
-    states = dict(mod.named_buffers())
-    states.update(mod.named_parameters())
+  """Returns a pytree of jax.ndarray and a jax callable."""
+  if env is None:
+    env = default_env()
+  states = dict(mod.named_buffers())
+  states.update(mod.named_parameters())
 
-    states = env.t2j_copy(states)
+  states = env.t2j_copy(states)
 
-    # @jax.jit
-    def jax_func(states, args, kwargs=None):
-        (states, args, kwargs) = env.j2t_iso((states, args, kwargs))
-        with env:
-            res = torch.func.functional_call(
-                mod, states, args, kwargs, tie_weights=False
-            )
-        return env.t2j_iso(res)
+  # @jax.jit
+  def jax_func(states, args, kwargs=None):
+    (states, args, kwargs) = env.j2t_iso((states, args, kwargs))
+    with env:
+      res = torch.func.functional_call(mod, states, args, kwargs, tie_weights=False)
+    return env.t2j_iso(res)
 
-    return states, jax_func
+  return states, jax_func
 
 
 def enable_globally():
-    env = default_env().enable_torch_modes()
-    return env
+  env = default_env().enable_torch_modes()
+  return env
 
 
 def disable_globally():
-    global env
-    default_env().disable_torch_modes()
+  global env
+  default_env().disable_torch_modes()
 
 
 @contextlib.contextmanager
 def disable_temporarily():
-    prev = default_env().enabled
-    if prev:
-        disable_globally()
-    yield ()
-    if prev:
-        enable_globally()
+  prev = default_env().enabled
+  if prev:
+    disable_globally()
+  yield ()
+  if prev:
+    enable_globally()
 
 
 torch.utils.rename_privateuse1_backend("jax")
 unsupported_dtype = [torch.quint8]
 
-import jax
-import torchax.device_module
 
 torch._register_device_module("jax", torchax.device_module)
 
 
 def enable_accuracy_mode():
-    jax.config.update("jax_enable_x64", True)
-    jax.config.update("jax_default_matmul_precision", "highest")
-    default_env().config.internal_respect_torch_return_dtypes = True
+  jax.config.update("jax_enable_x64", True)
+  jax.config.update("jax_default_matmul_precision", "highest")
+  default_env().config.internal_respect_torch_return_dtypes = True
 
 
 def enable_performance_mode():
-    jax.config.update("jax_enable_x64", False)
-    jax.config.update("jax_default_matmul_precision", "default")
-    default_env().config.internal_respect_torch_return_dtypes = False
+  jax.config.update("jax_enable_x64", False)
+  jax.config.update("jax_default_matmul_precision", "default")
+  default_env().config.internal_respect_torch_return_dtypes = False
 
 
 @dataclasses.dataclass
 class CompileOptions:
-    # only valid if compiling nn.Module
-    methods_to_compile: List[str] = dataclasses.field(
-        default_factory=lambda: ["forward"]
-    )
-    jax_jit_kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
-    mode: str = "jax"  # or dynamo or export
-
-
-def compile(fn, options: Optional[CompileOptions] = None):
-    options = options or CompileOptions()
-    if options.mode == "jax":
-        from torchax import interop
-
-        if isinstance(fn, torch.nn.Module):
-            module = interop.JittableModule(
-                fn, extra_jit_args=options.jax_jit_kwargs
-            )
-            for n in options.methods_to_compile:
-                module.make_jitted(n)
-            return module
-        else:
-            return interop.jax_jit(fn)
-    elif options.mode == "dynamo":
-        raise RuntimeError("dynamo mode is not supported yet")
-    elif options.mode == "export":
-        raise RuntimeError("export mode is not supported yet")
+  # only valid if compiling nn.Module
+  methods_to_compile: list[str] = dataclasses.field(default_factory=lambda: ["forward"])
+  jax_jit_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict)
+  mode: str = "jax"  # or dynamo or export
+
+
+def compile(fn, options: CompileOptions | None = None):
+  options = options or CompileOptions()
+  if options.mode == "jax":
+    from torchax import interop
+
+    if isinstance(fn, torch.nn.Module):
+      module = interop.JittableModule(fn, extra_jit_args=options.jax_jit_kwargs)
+      for n in options.methods_to_compile:
+        module.make_jitted(n)
+      return module
+    else:
+      return interop.jax_jit(fn)
+  elif options.mode == "dynamo":
+    raise RuntimeError("dynamo mode is not supported yet")
+  elif options.mode == "export":
+    raise RuntimeError("export mode is not supported yet")

torchax/amp.py

@@ -14,6 +14,7 @@
 
 import contextlib
 import enum
+
 import torch
 from torch.utils import _pytree as pytree
 
@@ -50,23 +51,24 @@ class CastPolicy(enum.Enum):
 
 
 def execute_policy(policy, args, kwargs, target_lower_fp):
-
   def is_float(a):
     return isinstance(a, torch.Tensor) and a.is_floating_point()
+
   match policy:
     case CastPolicy.LOWER_PRECISION_FP:
-      return pytree.tree_map_only(is_float, lambda a: a.to(target_lower_fp),
-                                  (args, kwargs))
+      return pytree.tree_map_only(
+        is_float, lambda a: a.to(target_lower_fp), (args, kwargs)
+      )
     case CastPolicy.FP32:
-      return pytree.tree_map_only(is_float, lambda a: a.to(torch.float32),
-                                  (args, kwargs))
+      return pytree.tree_map_only(
+        is_float, lambda a: a.to(torch.float32), (args, kwargs)
+      )
     case CastPolicy.PROMOTE:
-      dtypes = set(a.dtype for a in args)
+      dtypes = {a.dtype for a in args}
       widest = max((dtype.itemsize, dtype) for dtype in dtypes)[1]
-      return pytree.tree_map_only(is_float, lambda a: a.to(widest),
-                                  (args, kwargs))
+      return pytree.tree_map_only(is_float, lambda a: a.to(widest), (args, kwargs))
     case _:
-      raise AssertionError(f'Policy {policy} not implemented yet.')
+      raise AssertionError(f"Policy {policy} not implemented yet.")
 
 
 @contextlib.contextmanager
@@ -74,6 +76,7 @@ def autocast(device, dtype=torch.bfloat16, env=None):
   del device
   if env is None:
     import torchax
+
     env = torchax.default_env()
   with env.override_property(autocast_dtype=dtype):
     yield
@@ -81,266 +84,135 @@ def autocast(device, dtype=torch.bfloat16, env=None):
 
 # https://github.com/pytorch/pytorch/blob/05faba40287cf7d8734da96cb2e904f39710bf29/aten/src/ATen/autocast_mode.cpp#L327
 autocast_policy = {
-    torch.ops.aten.conv1d.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.conv1d.padding:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.conv2d.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.conv2d.padding:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.conv3d.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.conv3d.padding:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.bmm.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.mm.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.linalg_vecdot.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.baddbmm.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.addmm.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten._addmm_activation.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.addbmm.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.linear.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten._convolution.deprecated:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.matmul.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.conv_tbc.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.mkldnn_rnn_layer.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.conv_transpose1d.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.conv_transpose2d.input:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.conv_transpose3d.input:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.prelu.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten.scaled_dot_product_attention.default:
-        CastPolicy.LOWER_PRECISION_FP,
-    torch.ops.aten._native_multi_head_attention.default:
-        CastPolicy.LOWER_PRECISION_FP,
-
-    # fp32 cast policy
-    torch.ops.aten.avg_pool3d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.binary_cross_entropy.default:
-        CastPolicy.FP32,
-    torch.ops.aten.grid_sampler.default:
-        CastPolicy.FP32,
-    torch.ops.aten.polar.default:
-        CastPolicy.FP32,
-    torch.ops.aten.prod.default:
-        CastPolicy.FP32,
-    torch.ops.aten.prod.dim_int:
-        CastPolicy.FP32,
-    torch.ops.aten.prod.dim_Dimname:
-        CastPolicy.FP32,
-    torch.ops.aten.quantile.default:
-        CastPolicy.FP32,
-    torch.ops.aten.quantile.scalar:
-        CastPolicy.FP32,
-    torch.ops.aten.nanquantile.default:
-        CastPolicy.FP32,
-    torch.ops.aten.nanquantile.scalar:
-        CastPolicy.FP32,
-    torch.ops.aten.stft.default:
-        CastPolicy.FP32,
-    torch.ops.aten.stft.center:
-        CastPolicy.FP32,
-    torch.ops.aten.cdist.default:
-        CastPolicy.FP32,
-    torch.ops.aten.grid_sampler_2d.default:
-        CastPolicy.FP32,
-    torch.ops.aten._grid_sampler_2d_cpu_fallback.default:
-        CastPolicy.FP32,
-    torch.ops.aten.grid_sampler_3d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.trace.default:
-        CastPolicy.FP32,
-    torch.ops.aten.view_as_complex.default:
-        CastPolicy.FP32,
-    torch.ops.aten.cholesky.default:
-        CastPolicy.FP32,
-    torch.ops.aten.cholesky_inverse.default:
-        CastPolicy.FP32,
-    torch.ops.aten.cholesky_solve.default:
-        CastPolicy.FP32,
-    torch.ops.aten.inverse.default:
-        CastPolicy.FP32,
-    torch.ops.aten.lu_solve.default:
-        CastPolicy.FP32,
-    torch.ops.aten.orgqr.default:
-        CastPolicy.FP32,
-    torch.ops.aten.ormqr.default:
-        CastPolicy.FP32,
-    torch.ops.aten.pinverse.default:
-        CastPolicy.FP32,
-    torch.ops.aten.max_pool3d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.max_unpool2d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.max_unpool3d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.adaptive_avg_pool3d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.reflection_pad1d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.reflection_pad2d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.replication_pad1d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.replication_pad2d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.replication_pad3d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.mse_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.cosine_embedding_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.nll_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.nll_loss2d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.hinge_embedding_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.poisson_nll_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.smooth_l1_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.cross_entropy_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.l1_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.huber_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.margin_ranking_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.soft_margin_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.triplet_margin_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.multi_margin_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.ctc_loss.IntList:
-        CastPolicy.FP32,
-    torch.ops.aten.ctc_loss.Tensor:
-        CastPolicy.FP32,
-    torch.ops.aten.kl_div.default:
-        CastPolicy.FP32,
-    torch.ops.aten.multilabel_margin_loss.default:
-        CastPolicy.FP32,
-    torch.ops.aten.binary_cross_entropy_with_logits.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_fft.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_ifft.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_fft2.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_ifft2.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_fftn.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_ifftn.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_rfft.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_irfft.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_rfft2.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_irfft2.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_rfftn.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_irfftn.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_hfft.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fft_ihfft.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_cond.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_cond.p_str:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_matrix_rank.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_matrix_rank.tol_tensor:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_matrix_rank.atol_rtol_tensor:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_matrix_rank.atol_rtol_float:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_solve.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_cholesky.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_svdvals.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_eigvals.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_eigvalsh.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_inv.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_householder_product.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_tensorinv.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_tensorsolve.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fake_quantize_per_tensor_affine.default:
-        CastPolicy.FP32,
-    torch.ops.aten.geqrf.default:
-        CastPolicy.FP32,
-    torch.ops.aten._lu_with_info.default:
-        CastPolicy.FP32,
-    torch.ops.aten.qr.default:
-        CastPolicy.FP32,
-    torch.ops.aten.svd.default:
-        CastPolicy.FP32,
-    torch.ops.aten.triangular_solve.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fractional_max_pool2d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.fractional_max_pool3d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.adaptive_max_pool3d.default:
-        CastPolicy.FP32,
-    torch.ops.aten.multilabel_margin_loss_forward.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_qr.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_cholesky_ex.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_svd.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_eig.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_eigh.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_lstsq.default:
-        CastPolicy.FP32,
-    torch.ops.aten.linalg_inv_ex.default:
-        CastPolicy.FP32,
-
-    # promote
-    torch.ops.aten.stack.default:
-        CastPolicy.PROMOTE,
-    torch.ops.aten.cat.default:
-        CastPolicy.PROMOTE,
-    torch.ops.aten.index_copy.default:
-        CastPolicy.PROMOTE,
-    torch.ops.aten.index_copy.dimname:
-        CastPolicy.PROMOTE,
+  torch.ops.aten.conv1d.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv1d.padding: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv2d.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv2d.padding: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv3d.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv3d.padding: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.bmm.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.mm.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.linalg_vecdot.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.baddbmm.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.addmm.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten._addmm_activation.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.addbmm.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.linear.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten._convolution.deprecated: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.matmul.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv_tbc.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.mkldnn_rnn_layer.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv_transpose1d.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv_transpose2d.input: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv_transpose3d.input: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.prelu.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.scaled_dot_product_attention.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten._native_multi_head_attention.default: CastPolicy.LOWER_PRECISION_FP,
+  # fp32 cast policy
+  torch.ops.aten.avg_pool3d.default: CastPolicy.FP32,
+  torch.ops.aten.binary_cross_entropy.default: CastPolicy.FP32,
+  torch.ops.aten.grid_sampler.default: CastPolicy.FP32,
+  torch.ops.aten.polar.default: CastPolicy.FP32,
+  torch.ops.aten.prod.default: CastPolicy.FP32,
+  torch.ops.aten.prod.dim_int: CastPolicy.FP32,
+  torch.ops.aten.prod.dim_Dimname: CastPolicy.FP32,
+  torch.ops.aten.quantile.default: CastPolicy.FP32,
+  torch.ops.aten.quantile.scalar: CastPolicy.FP32,
+  torch.ops.aten.nanquantile.default: CastPolicy.FP32,
+  torch.ops.aten.nanquantile.scalar: CastPolicy.FP32,
+  torch.ops.aten.stft.default: CastPolicy.FP32,
+  torch.ops.aten.stft.center: CastPolicy.FP32,
+  torch.ops.aten.cdist.default: CastPolicy.FP32,
+  torch.ops.aten.grid_sampler_2d.default: CastPolicy.FP32,
+  torch.ops.aten._grid_sampler_2d_cpu_fallback.default: CastPolicy.FP32,
+  torch.ops.aten.grid_sampler_3d.default: CastPolicy.FP32,
+  torch.ops.aten.trace.default: CastPolicy.FP32,
+  torch.ops.aten.view_as_complex.default: CastPolicy.FP32,
+  torch.ops.aten.cholesky.default: CastPolicy.FP32,
+  torch.ops.aten.cholesky_inverse.default: CastPolicy.FP32,
+  torch.ops.aten.cholesky_solve.default: CastPolicy.FP32,
+  torch.ops.aten.inverse.default: CastPolicy.FP32,
+  torch.ops.aten.lu_solve.default: CastPolicy.FP32,
+  torch.ops.aten.orgqr.default: CastPolicy.FP32,
+  torch.ops.aten.ormqr.default: CastPolicy.FP32,
+  torch.ops.aten.pinverse.default: CastPolicy.FP32,
+  torch.ops.aten.max_pool3d.default: CastPolicy.FP32,
+  torch.ops.aten.max_unpool2d.default: CastPolicy.FP32,
+  torch.ops.aten.max_unpool3d.default: CastPolicy.FP32,
+  torch.ops.aten.adaptive_avg_pool3d.default: CastPolicy.FP32,
+  torch.ops.aten.reflection_pad1d.default: CastPolicy.FP32,
+  torch.ops.aten.reflection_pad2d.default: CastPolicy.FP32,
+  torch.ops.aten.replication_pad1d.default: CastPolicy.FP32,
+  torch.ops.aten.replication_pad2d.default: CastPolicy.FP32,
+  torch.ops.aten.replication_pad3d.default: CastPolicy.FP32,
+  torch.ops.aten.mse_loss.default: CastPolicy.FP32,
+  torch.ops.aten.cosine_embedding_loss.default: CastPolicy.FP32,
+  torch.ops.aten.nll_loss.default: CastPolicy.FP32,
+  torch.ops.aten.nll_loss2d.default: CastPolicy.FP32,
+  torch.ops.aten.hinge_embedding_loss.default: CastPolicy.FP32,
+  torch.ops.aten.poisson_nll_loss.default: CastPolicy.FP32,
+  torch.ops.aten.smooth_l1_loss.default: CastPolicy.FP32,
+  torch.ops.aten.cross_entropy_loss.default: CastPolicy.FP32,
+  torch.ops.aten.l1_loss.default: CastPolicy.FP32,
+  torch.ops.aten.huber_loss.default: CastPolicy.FP32,
+  torch.ops.aten.margin_ranking_loss.default: CastPolicy.FP32,
+  torch.ops.aten.soft_margin_loss.default: CastPolicy.FP32,
+  torch.ops.aten.triplet_margin_loss.default: CastPolicy.FP32,
+  torch.ops.aten.multi_margin_loss.default: CastPolicy.FP32,
+  torch.ops.aten.ctc_loss.IntList: CastPolicy.FP32,
+  torch.ops.aten.ctc_loss.Tensor: CastPolicy.FP32,
+  torch.ops.aten.kl_div.default: CastPolicy.FP32,
+  torch.ops.aten.multilabel_margin_loss.default: CastPolicy.FP32,
+  torch.ops.aten.binary_cross_entropy_with_logits.default: CastPolicy.FP32,
+  torch.ops.aten.fft_fft.default: CastPolicy.FP32,
+  torch.ops.aten.fft_ifft.default: CastPolicy.FP32,
+  torch.ops.aten.fft_fft2.default: CastPolicy.FP32,
+  torch.ops.aten.fft_ifft2.default: CastPolicy.FP32,
+  torch.ops.aten.fft_fftn.default: CastPolicy.FP32,
+  torch.ops.aten.fft_ifftn.default: CastPolicy.FP32,
+  torch.ops.aten.fft_rfft.default: CastPolicy.FP32,
+  torch.ops.aten.fft_irfft.default: CastPolicy.FP32,
+  torch.ops.aten.fft_rfft2.default: CastPolicy.FP32,
+  torch.ops.aten.fft_irfft2.default: CastPolicy.FP32,
+  torch.ops.aten.fft_rfftn.default: CastPolicy.FP32,
+  torch.ops.aten.fft_irfftn.default: CastPolicy.FP32,
+  torch.ops.aten.fft_hfft.default: CastPolicy.FP32,
+  torch.ops.aten.fft_ihfft.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_cond.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_cond.p_str: CastPolicy.FP32,
+  torch.ops.aten.linalg_matrix_rank.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_matrix_rank.tol_tensor: CastPolicy.FP32,
+  torch.ops.aten.linalg_matrix_rank.atol_rtol_tensor: CastPolicy.FP32,
+  torch.ops.aten.linalg_matrix_rank.atol_rtol_float: CastPolicy.FP32,
+  torch.ops.aten.linalg_solve.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_cholesky.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_svdvals.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_eigvals.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_eigvalsh.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_inv.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_householder_product.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_tensorinv.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_tensorsolve.default: CastPolicy.FP32,
+  torch.ops.aten.fake_quantize_per_tensor_affine.default: CastPolicy.FP32,
+  torch.ops.aten.geqrf.default: CastPolicy.FP32,
+  torch.ops.aten._lu_with_info.default: CastPolicy.FP32,
+  torch.ops.aten.qr.default: CastPolicy.FP32,
+  torch.ops.aten.svd.default: CastPolicy.FP32,
+  torch.ops.aten.triangular_solve.default: CastPolicy.FP32,
+  torch.ops.aten.fractional_max_pool2d.default: CastPolicy.FP32,
+  torch.ops.aten.fractional_max_pool3d.default: CastPolicy.FP32,
+  torch.ops.aten.adaptive_max_pool3d.default: CastPolicy.FP32,
+  torch.ops.aten.multilabel_margin_loss_forward.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_qr.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_cholesky_ex.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_svd.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_eig.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_eigh.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_lstsq.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_inv_ex.default: CastPolicy.FP32,
+  # promote
+  torch.ops.aten.stack.default: CastPolicy.PROMOTE,
+  torch.ops.aten.cat.default: CastPolicy.PROMOTE,
+  torch.ops.aten.index_copy.default: CastPolicy.PROMOTE,
+  torch.ops.aten.index_copy.dimname: CastPolicy.PROMOTE,
 }