From de355d278376bfc5379a4f9e598f757baf080d49 Mon Sep 17 00:00:00 2001 From: Nikita Shulga Date: Mon, 15 Jun 2026 11:36:05 -0700 Subject: [PATCH 1/3] Remove cholesky from FP32 cast policy `torch.cholesky` has been an alias for `torch.linalg.cholesky` for several years and is getting removed by https://github.com/pytorch/pytorch/pull/186817 Update the repo to use `torch.linalg.cholesky` --- torchax/amp.py | 1 - 1 file changed, 1 deletion(-) diff --git a/torchax/amp.py b/torchax/amp.py index ae5d5cf..2007dba 100644 --- a/torchax/amp.py +++ b/torchax/amp.py @@ -128,7 +128,6 @@ def autocast(device, dtype=torch.bfloat16, env=None): 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, From d067521ab6b13cf8e40d5a408eb8a187ab30665a Mon Sep 17 00:00:00 2001 From: Nikita Shulga Date: Mon, 15 Jun 2026 11:37:09 -0700 Subject: [PATCH 2/3] Rename cholesky operation to linalg_cholesky --- test/test_ops.py | 1 + torchax/ops/jaten.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_ops.py b/test/test_ops.py index 51f9f0f..a71f3f0 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -30,6 +30,7 @@ "_segment_reduce", "bincount", # NOTE: dtype for int input torch gives float. This is weird. "byte", + "cholesky", # To be removed from PyTorch, still covered by linalg.cholesky "cholesky_solve", "geqrf", "histogram", # hard op: AssertionError: Tensor-likes are not close! diff --git a/torchax/ops/jaten.py b/torchax/ops/jaten.py index e2d72d2..9bd3643 100644 --- a/torchax/ops/jaten.py +++ b/torchax/ops/jaten.py @@ -203,7 +203,7 @@ def _aten_index_select(x, dim, index): return jnp.take(x, index, dim) -@op(torch.ops.aten.cholesky) +@op(torch.ops.aten.linalg_cholesky) def _aten_cholesky(input, upper=False): return jax.scipy.linalg.cholesky(input, lower=(not upper)) From 9e573cb1afb156ddc891e8409c44949887350f13 Mon Sep 17 00:00:00 2001 From: Nikita Shulga Date: Thu, 2 Jul 2026 16:13:05 -0700 Subject: [PATCH 3/3] Fix lint --- test/test_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_ops.py b/test/test_ops.py index a71f3f0..4e72190 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -30,7 +30,7 @@ "_segment_reduce", "bincount", # NOTE: dtype for int input torch gives float. This is weird. "byte", - "cholesky", # To be removed from PyTorch, still covered by linalg.cholesky + "cholesky", # To be removed from PyTorch, still covered by linalg.cholesky "cholesky_solve", "geqrf", "histogram", # hard op: AssertionError: Tensor-likes are not close!