# Owner(s): ["module: nestedtensor"]

import io
import itertools
import math
import sys
import unittest
from functools import partial
from typing import Optional, Tuple

import numpy as np

import torch
import torch._dynamo
import torch._dynamo.testing
import torch.nn
import torch.nn.functional as F
from torch.nested._internal.nested_tensor import (
    buffer_from_jagged,
    jagged_from_list,
    nested_view_from_values_offsets,
    NestedTensor,
    ViewNestedFromBuffer,
)
from torch.testing._internal.common_cuda import (
    PLATFORM_SUPPORTS_FUSED_ATTENTION,
    SM70OrLater,
    SM80OrLater,
)
from torch.testing._internal.common_device_type import (
    dtypes,
    dtypesIfCUDA,
    instantiate_device_type_tests,
    onlyCPU,
    onlyCUDA,
    ops,
    PYTORCH_CUDA_MEMCHECK,
    skipCPUIf,
    skipCUDAIf,
    skipCUDAIfRocm,
    skipMeta,
)
from torch.testing._internal.common_dtype import floating_types_and_half
from torch.testing._internal.common_utils import (
    decorateIf,
    freeze_rng_state,
    gradcheck,
    instantiate_parametrized_tests,
    IS_FBCODE,
    IS_WINDOWS,
    markDynamoStrictTest,
    NestedTensorTestCase,
    parametrize,
    run_tests,
    skipIfSlowGradcheckEnv,
    skipIfTorchDynamo,
    subtest,
    TEST_WITH_ROCM,
    xfailIfTorchDynamo,
)
from torch.testing._internal.opinfo.definitions.nested import njt_op_db
from torch.utils._pytree import tree_flatten
from torch.utils.checkpoint import checkpoint, create_selective_checkpoint_contexts


# Tests are ported from pytorch/nestedtensor.
# This makes porting as_nested_tensor easier in the future.


def _iter_constructors():
    # yield as_nested_tensor
    yield torch.nested.nested_tensor


# Returns True if the function recompiles between inputs1 and inputs2 with the
# specified dynamic setting.
def _recompiles_for_inputs(fn, inputs1, inputs2, dynamic=True):
    compile_count = [0]

    def counter(gm, example_inputs):
        compile_count[0] += 1
        return gm

    compiled_f = torch.compile(fn, fullgraph=True, backend=counter, dynamic=dynamic)
    compiled_f(*inputs1)
    compiled_f(*inputs2)
    return compile_count[0] > 1


# Helper function to generate a pair of random nested tensors
# one is contiguous, the other is not, but they appear to have same entries
# an output nested tensor consists of
# * `len(ragged_sizes)` matrices
# * matrices[i].shape == (20, ragged_sizes[i])


def random_nt_noncontiguous_pair(ragged_sizes, device="cpu", dtype=torch.float16):
    xs = []
    for size in ragged_sizes:
        xs.append(torch.randn((size, 20), device=device, dtype=dtype))
    # contiguous nested tensor
    ys = []
    for x in xs:
        ys.append(x.transpose(-1, -2))
    nt_contiguous = torch.nested.nested_tensor(ys)
    # noncontiguous nested tensor
    n = len(ragged_sizes)
    nt_noncontiguous = torch.nested.nested_tensor(xs).transpose(-1, -2)
    return nt_contiguous, nt_noncontiguous


# Helper functions to pad a noncontiguous nested tensor
# can be replaced once to_padded_tensor supports noncontiguous memory


def noncontiguous_to_padded_tensor(input, shape=None):
    tensors = input.unbind()
    ntensors = len(tensors)
    assert ntensors > 0
    if shape is None:
        shape = []
        for size in tensors[0].shape:
            shape.append(size)
        for i in range(1, ntensors):
            new_shape = tensors[i].shape
            for j in range(len(shape)):
                shape[j] = max(shape[j], new_shape[j])
        shape = [ntensors] + shape
    result = tensors[0].new_zeros(shape)
    for itensor in range(ntensors):
        tensor = tensors[itensor]
        view = result[itensor]
        for idim in range(tensor.dim()):
            view = view.narrow(idim, 0, tensor.size(idim))
        view.copy_(tensor)
    return result


# Helper function to generate a random nested tensor


def random_nt(
    device,
    dtype,
    num_tensors,
    max_dims,
    min_dims=None,
    layout=torch.strided,
    require_non_empty=True,
):
    if min_dims is None:
        min_dims = tuple([0] * len(max_dims))

    assert len(max_dims) == len(min_dims)
    for min_dim, max_dim in zip(min_dims, max_dims):
        assert max_dim > min_dim, "random_nt: max_dim must be greater than min_dim"
        assert min_dim >= 0, "random_nt: min_dim must be non-negative"
        if require_non_empty:
            assert not (
                min_dim == 0 and max_dim == 1
            ), "random_nt: zero cannot be the only possible value if require_non_empty is True"

    if require_non_empty:
        # Select a random idx that will be required to be non-empty
        non_zero_idx = torch.randint(low=0, high=num_tensors, size=(1,)).item()

    ts1 = []
    for i, _ in enumerate(range(num_tensors)):
        tensor_dims = []
        for min_dim, max_dim in zip(min_dims, max_dims):
            new_min_dim = min_dim
            if require_non_empty and i == non_zero_idx and min_dim == 0:
                new_min_dim = 1
            tensor_dims.append(
                torch.randint(low=new_min_dim, high=max_dim, size=(1,)).item()
            )
        t1 = torch.randn(tensor_dims, device=device, dtype=dtype)
        ts1.append(t1)

    return torch.nested.nested_tensor(ts1, device=device, dtype=dtype, layout=layout)


# Alternate approach to generating a random NT.
# dims should be something like [5, None, 10], with None indicating that a
# random ragged structure should be used
def random_nt_from_dims(
    dims, device=None, dtype=None, layout=torch.strided, requires_grad=False
):
    sizes = [
        [
            d if d is not None else torch.randint(2, 10, size=(1,)).item()
            for d in dims[1:]
        ]
        for d in range(dims[0])
    ]
    return torch.nested.nested_tensor(
        [torch.randn(*size) for size in sizes],
        device=device,
        dtype=dtype,
        layout=layout,
        requires_grad=requires_grad,
    )


# Creates an NT matching another NT's number of components and
# shape / ragged structure for all dims specified to be -1.
def random_nt_from_similar(other, dims=None):
    if dims is None:
        return torch.randn_like(other)
    assert len(dims) == other.dim()
    assert dims[0] == -1 or dims[0] == other.size(0)

    ret_sizes = []
    for t in other.unbind():
        other_size = t.shape
        ret_size = []
        for i, d in enumerate(dims[1:]):
            if d == -1:
                ret_size.append(other_size[i])
            else:
                ret_size.append(d)
        ret_sizes.append(ret_size)

    return torch.nested.nested_tensor(
        [torch.randn(*size) for size in ret_sizes], device=other.device
    )


# makes naming nice for tests that parametrize over layout.
def layout_name(layout):
    # e.g. "torch.jagged" -> "jagged"
    return layout.__repr__().split(".")[-1]


def get_op_name(layout):
    # e.g. "<OpOverload(op='aten.sum', overload='dim_IntList')>" -> "sum"
    return layout.__name__.split(".")[0].split("_")[-1]


# Helper function for test_dummy_mha_with_nt
@torch.fx.wrap
def convert_dense_to_nested_tensor_legacy(values):
    offsets = torch.arange(
        0, values.shape[0] * values.shape[1] + 1, values.shape[1], device=values.device
    )
    metadata_cache = {"max_seqlen": values.shape[1], "min_seqlen": 1}
    nt = ViewNestedFromBuffer.apply(
        values.view(-1, values.shape[-1]), offsets, metadata_cache
    )
    return nt


# Helper function for test_dummy_mha_with_nt
@torch.fx.wrap
def convert_jagged_to_nested_tensor_legacy(
    values: torch.Tensor, offsets: torch.Tensor, max_length: int
) -> torch.Tensor:
    metadata_cache = {"max_seqlen": max_length, "min_seqlen": 1}
    nt = ViewNestedFromBuffer.apply(values, offsets, metadata_cache)
    return nt


# Helper function for test_dummy_mha_with_nt
@torch.fx.wrap
def convert_nt_to_jagged_legacy(nt):
    return buffer_from_jagged(nt)


# Helper function for test_dummy_mha_with_nt
@torch.fx.wrap
def convert_dense_to_nested_tensor(values):
    nt = torch.nested.as_nested_tensor(values, layout=torch.jagged)
    return nt


# Helper function for test_dummy_mha_with_nt
@torch.fx.wrap
def convert_jagged_to_nested_tensor(
    values: torch.Tensor, offsets: torch.Tensor, max_length: int
) -> torch.Tensor:
    nt = torch.nested.nested_tensor_from_jagged(
        values, offsets, lengths=None, min_seqlen=1, max_seqlen=max_length
    )
    return nt


# Helper function for test_dummy_mha_with_nt
def convert_nt_to_jagged(nt):
    return nt.values()


@markDynamoStrictTest
class TestNestedTensor(NestedTensorTestCase):
    @parametrize("batch_size", [2, 4])
    @parametrize("max_seq_len", [3, 5])
    @parametrize("vocab_size", [10, 20])
    def test_2d_nested_tensor(self, batch_size, max_seq_len, vocab_size):
        data = []
        nested_tensor_ref_list = []
        for _ in range(batch_size):
            if max_seq_len == 0:
                length = 0
            else:
                length = np.random.randint(low=1, high=max_seq_len)
            row = list(np.random.randint(low=0, high=vocab_size, size=(length,)))
            data.append(row)
            nested_tensor_ref_list.append(torch.Tensor(row))
        nested_tensor = torch.nested.nested_tensor(data, dtype=torch.int64)
        nested_tensor_list = nested_tensor.unbind()
        for id in range(batch_size):
            self.assertEqual(
                nested_tensor_list[id], nested_tensor_ref_list[id].type(torch.int64)
            )

    @parametrize("batch_size", [2, 4])
    @parametrize("max_seq_len", [3, 5])
    @parametrize("vocab_size", [10, 20])
    def test_3d_nested_tensor(self, batch_size, max_seq_len, vocab_size):
        data = []
        nested_tensor_ref_list = []
        for _ in range(batch_size):
            if max_seq_len == 0:
                length = 0
            else:
                length = np.random.randint(low=1, high=max_seq_len)
            row = list(np.random.randint(low=0, high=vocab_size, size=(length,)))
            row = [list(item * np.arange(max_seq_len)) for item in row]
            data.append(row)
            nested_tensor_ref_list.append(torch.Tensor(row))
        nested_tensor = torch.nested.nested_tensor(data, dtype=torch.int64)
        nested_tensor_list = nested_tensor.unbind()
        for id in range(batch_size):
            self.assertEqual(
                nested_tensor_list[id], nested_tensor_ref_list[id].type(torch.int64)
            )

    @parametrize("batch_size", [2, 4])
    @parametrize("max_seq_len", [3, 5])
    @parametrize("vocab_size", [10, 20])
    def test_3d_nested_tensor_float(self, batch_size, max_seq_len, vocab_size):
        data = []
        nested_tensor_ref_list = []
        for _ in range(batch_size):
            if max_seq_len == 0:
                length = 0
            else:
                length = np.random.randint(low=1, high=max_seq_len)
            row = list(
                np.random.randint(low=0, high=vocab_size, size=(length,)).astype(float)
            )
            row = [list(item * np.arange(max_seq_len)) for item in row]
            data.append(row)
            nested_tensor_ref_list.append(torch.Tensor(row))
        nested_tensor = torch.nested.nested_tensor(data, dtype=torch.float)
        nested_tensor_list = nested_tensor.unbind()
        for id in range(batch_size):
            self.assertEqual(
                nested_tensor_list[id], nested_tensor_ref_list[id].type(torch.float)
            )

    @torch.inference_mode()
    def _test_unbind_case(self, a, b):
        nt = torch.nested.nested_tensor([a, b])
        a1, b1 = nt.unbind()
        self.assertTrue(a is not a1)
        self.assertTrue(b is not b1)

        nt = torch.nested.nested_tensor([a, b], dtype=a.dtype)
        a1, b1 = nt.unbind(0)
        self.assertEqual(a, a1)
        self.assertEqual(b, b1)

        a = torch.randn((2, 3)).add_(1)
        nt = torch.nested.nested_tensor([a])
        self.assertEqual(a, nt.unbind(0)[0])

    @torch.inference_mode()
    def test_unbind_0(self):
        self._test_unbind_case(torch.tensor([1, 2]), torch.tensor([7, 8]))

    @torch.inference_mode()
    def test_unbind_1(self):
        self._test_unbind_case(torch.tensor([1]), torch.tensor([7]))

    @torch.inference_mode()
    def test_unbind_3(self):
        self._test_unbind_case(torch.tensor([1.0]), torch.tensor([]))

    @torch.inference_mode()
    def test_unbind_4(self):
        self._test_unbind_case(torch.tensor([]), torch.tensor([]))

    @torch.inference_mode()
    def test_unbind_dim(self):
        def _test_fn(unbind_fn):
            a = torch.rand(3, 2)
            b = torch.rand(2, 3)
            nt = torch.nested.nested_tensor([a, b])
            self.assertRaises(RuntimeError, lambda: unbind_fn(nt, 1))

        # Both of these tests are necessary, because we're using
        # torch_function.
        _test_fn(lambda x, dim: x.unbind(dim))
        # TODO: Re-enable this once using torch_dispatch
        # _test_fn(lambda x, dim: torch.unbind(x, dim))

    @torch.inference_mode()
    def test_nested_tensor(self):
        self.assertRaises(
            TypeError, lambda: torch.nested.nested_tensor(torch.tensor([3.0]))
        )
        self.assertRaises(TypeError, lambda: torch.nested.nested_tensor(4.0))

    @torch.inference_mode()
    def test_nested_tensor_matching_dim(self):
        self.assertRaisesRegex(
            RuntimeError,
            "Found dimension 1 for Tensor at index 1 and dimension 0 for Tensor at index 0.",
            lambda: torch.nested.nested_tensor([torch.tensor(1.0), torch.tensor([])]),
        )
        self.assertRaisesRegex(
            RuntimeError,
            "Found dimension 1 for Tensor at index 2 and dimension 0 for Tensor at index 1.",
            lambda: torch.nested.nested_tensor(
                [torch.tensor(1.0), torch.tensor(2.0), torch.tensor([])]
            ),
        )

    @torch.inference_mode()
    def test_default_nested_tensor(self):
        self.assertRaises(TypeError, lambda: torch.nested.nested_tensor())
        default_nested_tensor = torch.nested.nested_tensor([])
        default_tensor = torch.tensor([])
        # self.assertEqual(default_nested_tensor.nested_dim(), 1)
        # self.assertEqual(default_nested_tensor.nested_size(), ())
        self.assertEqual(default_nested_tensor.dim(), default_tensor.dim())
        self.assertEqual(default_nested_tensor.layout, default_tensor.layout)
        self.assertEqual(default_nested_tensor.device, default_tensor.device)
        self.assertEqual(default_nested_tensor.dtype, default_tensor.dtype)
        self.assertEqual(
            default_nested_tensor.requires_grad, default_tensor.requires_grad
        )
        self.assertIsNone(default_tensor.grad)
        # TODO: Re-enable once we have a performance driven
        # use case and implementation.
        # self.assertEqual(default_nested_tensor.is_pinned(),
        #                  default_tensor.is_pinned())

    @torch.inference_mode()
    def test_dim(self):
        for constructor in _iter_constructors():
            a1 = constructor([])
            self.assertEqual(a1.dim(), 1)
            a1 = constructor([torch.tensor(3.0)])
            self.assertEqual(a1.dim(), 1)
            a1 = constructor([torch.tensor([1, 2, 3, 4])])
            self.assertEqual(a1.dim(), 2)

    @unittest.skipIf(IS_FBCODE, "numel is not virtual in fbcode.")
    @torch.inference_mode()
    def test_numel(self):
        for constructor in _iter_constructors():
            a1 = constructor([])
            self.assertEqual(a1.numel(), 0)
            a1 = constructor([torch.tensor(3.0), torch.tensor(4.0)])
            self.assertEqual(a1.numel(), 2)
            a1 = constructor([torch.randn(2, 2, 2)])
            self.assertEqual(a1.numel(), 8)
            a1 = constructor([torch.randn([1, 2, 3]), torch.randn(3, 2, 1)])
            self.assertEqual(a1.numel(), 12)
            a1 = constructor([torch.randn([1, 1, 3]), torch.randn(3, 2, 4)])
            self.assertEqual(a1.numel(), 27)
            a1 = constructor([torch.randn([5, 5, 5]), torch.randn(6, 6, 6)])
            self.assertEqual(a1.numel(), 341)

            # Interesting edge case
            a1 = constructor([torch.randn([1, 2, 3]), torch.randn(1, 2, 0)])
            self.assertEqual(a1.numel(), 6)

    @torch.inference_mode()
    def test_size(self):
        for constructor in _iter_constructors():
            a1 = constructor([])
            self.assertRaisesRegex(
                RuntimeError,
                "NestedTensorImpl doesn't support sizes",
                lambda: a1.size(),
            )

    def test_size_dim(self):
        a = torch.nested.nested_tensor([])
        self.assertEqual(a.size(0), 0)

        a = torch.nested.nested_tensor([torch.tensor(1)])
        self.assertEqual(a.size(0), 1)

        a = torch.nested.nested_tensor([torch.tensor(1), torch.tensor(2)])
        self.assertEqual(a.size(0), 2)

        a = torch.nested.nested_tensor([torch.rand(1, 2), torch.rand(1, 8)])
        self.assertEqual(a.size(0), 2)
        self.assertEqual(a.size(1), 1)
        self.assertRaisesRegex(
            RuntimeError,
            "Given dimension 2 is irregular and does not have a size",
            lambda: a.size(2),
        )

        a = torch.nested.nested_tensor([torch.rand(3, 4), torch.rand(5, 4)])
        self.assertEqual(a.size(0), 2)
        self.assertRaisesRegex(
            RuntimeError,
            "Given dimension 1 is irregular and does not have a size",
            lambda: a.size(1),
        )
        self.assertEqual(a.size(2), 4)

    @unittest.skipIf(IS_FBCODE, "stride is not virtual in fbcode.")
    @torch.inference_mode()
    def test_stride(self):
        for constructor in _iter_constructors():
            a1 = constructor([])
            self.assertRaisesRegex(
                RuntimeError,
                "NestedTensorImpl doesn't support strides",
                lambda: a1.stride(),
            )

    @unittest.skipIf(IS_FBCODE, "is_contiguous is not virtual in fbcode.")
    @torch.inference_mode()
    def test_is_contiguous(self):
        # Test empty case
        nt_empty = torch.nested.nested_tensor([])
        assert nt_empty.is_contiguous()
        self.assertEqual(nt_empty, nt_empty.contiguous())

        nt_contiguous, nt_noncontiguous = random_nt_noncontiguous_pair((2, 3, 6, 7))

        # Test contiguous case
        assert nt_contiguous.is_contiguous()
        self.assertEqual(nt_contiguous, nt_contiguous.contiguous())

        # Test non_contiguous case
        assert not nt_noncontiguous.is_contiguous()
        self.assertEqual(nt_contiguous, nt_noncontiguous.contiguous())

        # Test querying by memory_format
        self.assertTrue(
            nt_contiguous.is_contiguous(memory_format=torch.contiguous_format)
        )
        self.assertTrue(
            not nt_noncontiguous.is_contiguous(memory_format=torch.contiguous_format)
        )

    @torch.inference_mode()
    def test_repr_string(self):
        a = torch.nested.nested_tensor([])
        expected = "nested_tensor([\n\n])"
        self.assertEqual(str(a), expected)
        self.assertEqual(repr(a), expected)

        a = torch.nested.nested_tensor([torch.tensor(1.0)])
        expected = "nested_tensor([\n  tensor(1.)\n])"
        self.assertEqual(str(a), expected)
        self.assertEqual(repr(a), expected)

        a = torch.nested.nested_tensor([torch.tensor([[1, 2]]), torch.tensor([[4, 5]])])
        expected = "nested_tensor([\n  tensor([[1, 2]]),\n  tensor([[4, 5]])\n])"
        self.assertEqual(str(a), expected)
        self.assertEqual(repr(a), expected)

    def test_to_padded_tensor_on_empty_tensor(self):
        nt = torch.nested.nested_tensor([])
        empty = torch.nested.to_padded_tensor(nt, 4)
        self.assertEqual(empty, torch.tensor([]))

    def test_nested_namespace(self):
        nt = torch.nested.nested_tensor([torch.randn(2, 3), torch.randn(4, 5)])
        result = nt.to_padded_tensor(4)
        nested_namespace_result = torch.nested.to_padded_tensor(nt, 4)
        self.assertEqual(result, nested_namespace_result)

    def test_to(self):
        ntensors = 4
        nt = random_nt(torch.device("cpu"), torch.float32, ntensors, (4, 4))

        def test_copy_behavior(t, non_blocking=False):
            self.assertIs(t, t.to(t, non_blocking=non_blocking))
            self.assertIs(t, t.to(t.dtype, non_blocking=non_blocking))
            self.assertIs(t, t.to(torch.empty_like(t), non_blocking=non_blocking))
            self.assertIsNot(t, t.to(t, non_blocking=non_blocking, copy=True))
            self.assertIsNot(t, t.to(t.dtype, non_blocking=non_blocking, copy=True))
            self.assertIsNot(
                t, t.to(torch.empty_like(t), non_blocking=non_blocking, copy=True)
            )

            devices = [t.device]
            if t.device.type == "cuda":
                if t.device.index == -1:
                    devices.append(f"cuda:{torch.cuda.current_device()}")
                elif t.device.index == torch.cuda.current_device():
                    devices.append("cuda")
            for device in devices:
                self.assertIs(t, t.to(device, non_blocking=non_blocking))
                self.assertIs(t, t.to(device, t.dtype, non_blocking=non_blocking))
                self.assertIsNot(t, t.to(device, non_blocking=non_blocking, copy=True))
                self.assertIsNot(
                    t, t.to(device, t.dtype, non_blocking=non_blocking, copy=True)
                )

        test_copy_behavior(nt)
        self.assertEqual(nt.device, nt.to("cpu").device)
        self.assertEqual(nt.device, nt.to("cpu", dtype=torch.float32).device)
        self.assertIs(torch.float32, nt.to("cpu", dtype=torch.float32).dtype)
        self.assertEqual(nt.device, nt.to(torch.float32).device)
        self.assertIs(torch.float32, nt.to(dtype=torch.float32).dtype)

        def test_data_ptr(getter):
            self.assertEqual(getter(nt), getter(nt.to("cpu")))
            self.assertEqual(
                getter(nt), getter(nt.to(dtype=nt.dtype, device=nt.device, copy=False))
            )
            self.assertEqual(getter(nt), getter(nt.to("cpu", copy=False)))
            self.assertNotEqual(getter(nt), getter(nt.to("cpu", copy=True)))

        test_data_ptr(lambda nt: nt.data_ptr())

        if torch.cuda.is_available():
            for non_blocking in [True, False]:
                for cuda in [
                    "cuda",
                    "cuda:0" if torch.cuda.device_count() == 1 else "cuda:1",
                ]:
                    nt2 = random_nt(cuda, torch.float32, ntensors, (4, 4))
                    test_copy_behavior(nt2, non_blocking)
                    self.assertEqual(
                        nt2.device, nt2.to(cuda, non_blocking=non_blocking).device
                    )
                    self.assertEqual(
                        nt.device, nt2.to("cpu", non_blocking=non_blocking).device
                    )
                    self.assertEqual(
                        nt2.device, nt.to(cuda, non_blocking=non_blocking).device
                    )
                    self.assertIs(
                        torch.int32,
                        nt2.to(
                            "cpu", dtype=torch.int32, non_blocking=non_blocking
                        ).dtype,
                    )
                    self.assertEqual(
                        nt.device,
                        nt2.to(
                            "cpu", dtype=torch.int32, non_blocking=non_blocking
                        ).device,
                    )
                    self.assertIs(torch.int32, nt2.to(dtype=torch.int32).dtype)
                    self.assertEqual(nt2.device, nt2.to(dtype=torch.int32).device)

    def test_copy_(self):
        ntensors = 4
        nt = random_nt(torch.device("cpu"), torch.float32, ntensors, (4, 4))
        nt_copy = torch.empty_like(nt)
        nt_copy.copy_(nt)

        for nt_ub, nt_copy_ub in zip(nt.unbind(), nt_copy):
            self.assertEqual(nt_ub, nt_copy_ub)

        nt_error = torch.nested.nested_tensor([torch.tensor([0, 0])])
        self.assertRaisesRegex(
            RuntimeError,
            "copy_ only supports tensors that are the same size for Nested implementations",
            lambda: nt_error.copy_(nt),
        )

        if torch.cuda.is_available():
            nt = random_nt(torch.device("cuda"), torch.float32, ntensors, (4, 4))
            nt_copy = torch.empty_like(nt, device=torch.device("cpu"))
            nt_copy.copy_(nt, non_blocking=True)
            torch.cuda.current_stream(torch.cuda.current_device()).synchronize()
            for nt_ub, nt_copy_ub in zip(nt.unbind(), nt_copy):
                self.assertEqual(nt_ub, nt_copy_ub)

            nt_copy = torch.empty_like(nt, device=torch.device("cpu"))
            nt_copy.copy_(nt, non_blocking=False)
            for nt_ub, nt_copy_ub in zip(nt.unbind(), nt_copy):
                self.assertEqual(nt_ub, nt_copy_ub)

    def test_fill_(self):
        ntensors = 4
        nt = random_nt(torch.device("cpu"), torch.float32, ntensors, (4, 4))
        nt.fill_(10.0)
        for nt_ub in nt.unbind():
            t = torch.empty_like(nt_ub)
            t.fill_(10.0)
            self.assertEqual(nt_ub, t)

        fill_tensor = torch.tensor([11.0])
        self.assertRaisesRegex(
            RuntimeError,
            "fill_ only supports 0-dimension value tensor",
            lambda: nt.fill_(fill_tensor),
        )

        nt.fill_(fill_tensor[0])
        for nt_ub in nt.unbind():
            t = torch.empty_like(nt_ub)
            t.fill_(11.0)
            self.assertEqual(nt_ub, t)

    def test_zero_(self):
        ntensors = 4
        nt = random_nt(torch.device("cpu"), torch.float32, ntensors, (4, 4))
        nt.zero_()
        for nt_ub in nt.unbind():
            t = torch.empty_like(nt_ub)
            t.fill_(0.0)
            self.assertEqual(nt_ub, t)

    @parametrize(
        "func",
        [torch.ones_like, torch.zeros_like, torch.randn_like],
        name_fn=lambda f: f.__name__,
    )
    def test_like_functions(self, func):
        ntensors = 4
        nt = random_nt(torch.device("cpu"), torch.float32, ntensors, (4, 4))
        torch.manual_seed(1)
        nt_like = func(nt)

        torch.manual_seed(1)
        for nt_ub in nt_like.unbind():
            t_like = func(nt_ub)
            self.assertEqual(nt_ub, t_like)

    def test_cat(self):
        # dim=0 success case
        # No constraints on ragged structures matching.
        x = random_nt_from_dims([5, None, 10])
        y = random_nt_from_dims([3, 4, None])
        output = torch.cat([x, y], dim=0)
        for out_component, xy_component in zip(
            output.unbind(), itertools.chain(x.unbind(), y.unbind())
        ):
            self.assertEqual(out_component, xy_component)

        # dim=-1 success case
        # shape (B, *, D)
        x = random_nt_from_dims([5, None, 10])
        # shape (B, *, D'); same structure as x but dim=-1 differs
        y = random_nt_from_similar(x, dims=[-1, -1, 8])
        # should be shape (B, *, D + D') when supported
        output = torch.cat([x, y], dim=-1)
        for out_component, x_component, y_component in zip(
            output.unbind(), x.unbind(), y.unbind()
        ):
            self.assertEqual(
                out_component, torch.cat([x_component, y_component], dim=-1)
            )

        # dim between 0 and -1 success case
        x = random_nt_from_dims([5, None, 2, 3])
        # same structure as x but dim=2 differs
        y = random_nt_from_similar(x, dims=[-1, -1, 4, -1])
        output = torch.cat([x, y], dim=2)
        for out_component, x_component, y_component in zip(
            output.unbind(), x.unbind(), y.unbind()
        ):
            self.assertEqual(
                out_component, torch.cat([x_component, y_component], dim=1)
            )

        # error case: mixed NT / dense inputs
        x = random_nt_from_dims([5, None, 2])
        y = torch.randn(5, 3, 2)
        with self.assertRaisesRegex(
            RuntimeError, "expected each tensor in given list to be nested"
        ):
            torch.cat([x, y], dim=-1)

        # error case: NTs with different dims
        x = random_nt_from_dims([5, None, 2])
        y = random_nt_from_dims([5, None, 2, 3])
        with self.assertRaisesRegex(
            RuntimeError,
            "expected all nested tensors to have matching ragged structures outside of the concatenated dim",
        ):
            torch.cat([x, y], dim=-1)

        # error case: non-contiguous NT
        x, y = random_nt_noncontiguous_pair((2, 3, 4), dtype=torch.float32)
        # transpose to put ragged dim next to batch dim
        x, y = x.transpose(-2, -1), y.transpose(-2, -1)
        with self.assertRaisesRegex(
            RuntimeError, "only contiguous nested tensors are supported"
        ):
            torch.cat([x, y], dim=-1)

        # error case: multiple ragged dims in inputs
        x = random_nt_from_dims([5, None, None, 2])
        y = random_nt_from_similar(x)
        with self.assertRaisesRegex(
            RuntimeError,
            "only nested tensors with a single ragged dim next to the batch dim are supported",
        ):
            torch.cat([x, y], dim=-1)

        # error case: ragged dim not next to batch dim
        x = random_nt_from_dims([5, 2, None])
        y = random_nt_from_similar(x)
        with self.assertRaisesRegex(
            RuntimeError,
            "only nested tensors with a single ragged dim next to the batch dim are supported",
        ):
            torch.cat([x, y], dim=1)

        # error case: NTs with different batch sizes
        x = random_nt_from_dims([5, None, 2])
        y = random_nt_from_dims([3, None, 2])
        with self.assertRaisesRegex(
            RuntimeError,
            "expected all nested tensors to have matching ragged structures outside of the concatenated dim",
        ):
            torch.cat([x, y], dim=-1)

        # error case: NTs with different ragged structures
        x = torch.nested.nested_tensor(
            [
                torch.randn(2, 6),
                torch.randn(4, 6),
                torch.randn(5, 6),
            ]
        )
        y = torch.nested.nested_tensor(
            [
                torch.randn(5, 6),
                torch.randn(4, 6),
                torch.randn(2, 6),
            ]
        )
        with self.assertRaisesRegex(
            RuntimeError,
            "expected all nested tensors to have matching ragged structures outside of the concatenated dim",
        ):
            torch.cat([x, y], dim=-1)


@markDynamoStrictTest
class TestNestedTensorDeviceType(NestedTensorTestCase):
    # Helper function to generate a pair of random nested tensors
    # the 2 nested tensors have same shapes
    def random_nt_pair(self, device, dtype, num_tensors, max_dims):
        ts1 = []
        ts2 = []
        for _ in range(num_tensors):
            tensor_dims = tuple(
                [
                    torch.randint(low=0, high=max_dim, size=(1,)).item()
                    for max_dim in max_dims
                ]
            )
            t1 = torch.randn(tensor_dims, device=device, dtype=dtype)
            t2 = torch.randn(tensor_dims, device=device, dtype=dtype)
            ts1.append(t1)
            ts2.append(t2)
        return (
            torch.nested.nested_tensor(ts1, device=device, dtype=dtype),
            torch.nested.nested_tensor(ts2, device=device, dtype=dtype),
        )

    @dtypes(*floating_types_and_half())
    def test_detach(self, device, dtype):
        a = torch.randn(2, 4, device=device, dtype=dtype, requires_grad=False)
        b = torch.randn(5, 4, device=device, dtype=dtype, requires_grad=False)
        x = torch.nested.nested_tensor([a, b], requires_grad=True)

        x_detach = x.detach()

        z = x_detach * 4
        self.assertFalse(x_detach.requires_grad)
        self.assertFalse(z.requires_grad)

        a = torch.randn(2, 4, device=device, dtype=dtype, requires_grad=True)
        b = torch.randn(5, 4, device=device, dtype=dtype, requires_grad=True)
        x = torch.nested.as_nested_tensor([a, b])

        y = x * 2
        y = y.detach()
        self.assertFalse(y.requires_grad)
        self.assertIsNone(y.grad_fn)

        z = x + y
        torch.nested.to_padded_tensor(z, 0).sum().backward()
        # This is an incorrect gradient, but we assume that's what the user
        # wanted. detach() is an advanced option.
        self.assertEqual(a.grad, torch.ones(2, 4, device=device, dtype=dtype))
        self.assertEqual(b.grad, torch.ones(5, 4, device=device, dtype=dtype))

    @dtypes(torch.float, torch.float16, torch.double)
    def test_unbind_noncontiguous(self, device, dtype):
        nt_contiguous, nt_noncontiguous = random_nt_noncontiguous_pair(
            (2, 3, 6, 7), device, dtype
        )
        ub_contiguous = nt_contiguous.unbind()
        ub_noncontiguous = nt_noncontiguous.unbind()
        self.assertEqual(len(ub_contiguous), len(ub_noncontiguous))
        n = len(ub_contiguous)
        for i in range(n):
            self.assertEqual(ub_contiguous[i], ub_noncontiguous[i])

    @dtypes(torch.float)
    @skipMeta
    def test_to_then_from_padded_tensor_no_transform0213(self, device, dtype):
        t = torch.randn(4, 4, 4, device=device, dtype=dtype)
        ts = list(torch.unbind(t))
        ts[0] = ts[0][:-1]
        nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
        padded = torch.nested.to_padded_tensor(nt, 0)

        nt_to = torch._nested_from_padded_and_nested_example(padded, nt)

        for t1, t2 in zip(nt.unbind(), nt_to.unbind()):
            self.assertEqual(t1, t2)
        self.assertEqual(nt.device, nt_to.device)

    @dtypes(torch.float)
    @dtypesIfCUDA(torch.float, torch.half)
    @skipMeta
    @torch.inference_mode()
    def test_layer_norm(self, device, dtype):
        def _test(size):
            # Simple shapes test
            t0 = torch.randn(2, size, device=device, dtype=dtype, requires_grad=False)
            t1 = torch.randn(2, size, device=device, dtype=dtype, requires_grad=False)
            ts = [t0, t1, t0, t1]
            nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
            layer_norm = torch.nn.LayerNorm(size, device=device, dtype=dtype)
            nt_result = layer_norm(nt)
            for nt_subresult, t in zip(nt_result.unbind(), ts):
                t_result = layer_norm(t.reshape(1, -1, size).squeeze(0))
                self.assertEqual(nt_subresult, t_result)

            # More complex nt test with different lengths for each tensor
            t0 = torch.randn(4, size, device=device, dtype=dtype, requires_grad=False)
            t1 = torch.randn(10, size, device=device, dtype=dtype, requires_grad=False)
            t2 = torch.randn(7, size, device=device, dtype=dtype, requires_grad=False)
            ts = [t0, t1, t2, t0, t2]
            nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
            layer_norm = torch.nn.LayerNorm(size, device=device, dtype=dtype)
            nt_result = layer_norm(nt)
            for nt_subresult, t in zip(nt_result.unbind(), ts):
                t_result = layer_norm(t.reshape(1, -1, size).squeeze(0))
                self.assertEqual(nt_subresult, t_result)

            if size <= 128:
                # Test with multidimensional tensors after irregular dim
                # (run only with smaller dimensions to ensure fast execution)
                t0 = torch.randn(
                    4, size, size, 4, device=device, dtype=dtype, requires_grad=False
                )
                t1 = torch.randn(
                    10, size, size, 4, device=device, dtype=dtype, requires_grad=False
                )
                t2 = torch.randn(
                    7, size, size, 4, device=device, dtype=dtype, requires_grad=False
                )
                ts = [t0, t1, t2, t0, t2]
                nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
                layer_norm = torch.nn.LayerNorm(
                    (size, size, 4), device=device, dtype=dtype
                )
                nt_result = layer_norm(nt)
                for nt_subresult, t in zip(nt_result.unbind(), ts):
                    t_result = layer_norm(t.reshape(1, -1, size, size, 4).squeeze(0))
                    self.assertEqual(nt_subresult, t_result)

                # Test where the normalizing dimensions are not all
                layer_norm = torch.nn.LayerNorm((size, 4), device=device, dtype=dtype)
                nt_result = layer_norm(nt)
                for nt_subresult, t in zip(nt_result.unbind(), ts):
                    t_result = layer_norm(t.reshape(1, -1, size, size, 4).squeeze(0))
                    self.assertEqual(nt_subresult, t_result)

        for size in (1024, 1023, 513, 512, 256, 128, 2, 4, 32):
            _test(size)

    @dtypes(torch.float)
    @dtypesIfCUDA(torch.float, torch.half)
    @skipMeta
    @torch.inference_mode()
    def test_layer_norm_breaking(self, device, dtype):
        size = 128
        t0 = torch.randn(
            4, size, size, 4, device=device, dtype=dtype, requires_grad=False
        )
        t1 = torch.randn(
            10, size, size, 4, device=device, dtype=dtype, requires_grad=False
        )
        t2 = torch.randn(
            7, size, size, 4, device=device, dtype=dtype, requires_grad=False
        )
        ts = [t0, t1, t2, t0, t2]
        nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
        layer_norm = torch.nn.LayerNorm((4, size, size, 4), device=device, dtype=dtype)
        self.assertRaisesRegex(
            RuntimeError,
            "normalized_shape extends into irregular dimensions for the nested tensor",
            lambda: layer_norm(nt),
        )
        layer_norm = torch.nn.LayerNorm((size + 1, size, 4), device=device, dtype=dtype)
        self.assertRaisesRegex(
            RuntimeError,
            "The shape at dimension 0",
            lambda: layer_norm(nt),
        )

    @parametrize("layout", [torch.strided, torch.jagged], name_fn=layout_name)
    def test_embedding(self, device, layout):
        inputs = [
            torch.randint(100, (L,), device=device, dtype=torch.int64)
            for L in torch.randint(5, 50, (8,))
        ]
        x = torch.nested.nested_tensor(
            inputs, device=device, dtype=torch.int64, layout=layout
        )
        emb = torch.nn.Embedding(100, 8, device=device)
        y = emb(x)

        @torch._dynamo.disable
        def check(inputs, y):
            ys = y.unbind()
            for i, inp in enumerate(inputs):
                self.assertEqual(emb(inp), ys[i])

        check(inputs, y)

    @skipMeta
    @torch.inference_mode()
    @dtypes(*floating_types_and_half())
    def test_masked_fill(self, device, dtype):
        # nested tensor * nested tensor
        (nt, mask) = self.random_nt_pair(device, dtype, 4, (4, 4))
        mask = torch.nested.nested_tensor([m < 0 for m in mask.unbind()])
        ref = torch.nested.nested_tensor(
            [t.masked_fill(m, 0) for (t, m) in zip(nt.unbind(), mask.unbind())]
        )
        out = nt.masked_fill(mask, 0)
        self.assertEqual(ref, out)

    @dtypes(torch.float, torch.float16)
    def test_to_padded_tensor_simple(self, device, dtype):
        t = torch.randn(4, 4, 4, device=device, dtype=dtype)
        ts = list(torch.unbind(t))
        ts[0] = ts[0][:-1]
        nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
        for padding_value in (0, 1):
            padded = torch.nested.to_padded_tensor(nt, padding_value)

            correct_output = t.clone()
            if padding_value == 0:
                correct_output[0][-1] = torch.zeros_like(correct_output[0][-1])
            else:
                correct_output[0][-1] = torch.ones_like(correct_output[0][-1])

            self.assertEqual(padded, correct_output)
            self.assertEqual(padded.device, torch.device(device))
            self.assertEqual(padded.dtype, dtype)

    @dtypes(torch.float, torch.float16)
    def test_to_padded_tensor_output_size(self, device, dtype):
        t = torch.randn(4, 4, 4, device=device, dtype=dtype)
        output_size = (4, 6, 5)
        ts = list(torch.unbind(t))
        ts[0] = ts[0][:-1]
        nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
        for padding_value in (0, 1):
            padded = torch.nested.to_padded_tensor(
                nt, padding_value, output_size=output_size
            )
            correct_output = (
                torch.ones(output_size, device=device, dtype=dtype) * padding_value
            )
            correct_output[:4:, :4, :4] = t.clone()
            if padding_value == 0:
                correct_output[0][3] = torch.zeros_like(correct_output[0][3])
            else:
                correct_output[0][3] = torch.ones_like(correct_output[0][3])

            self.assertEqual(padded, correct_output)
            self.assertEqual(padded.device, torch.device(device))
            self.assertEqual(padded.dtype, dtype)

    @dtypes(torch.float, torch.float16, torch.double)
    def test_to_padded_tensor_dim2(self, device, dtype):
        ts = [
            torch.randn(160, device=device, dtype=dtype),
            torch.randn(1240, device=device, dtype=dtype),
            torch.randn(2400, device=device, dtype=dtype),
        ]
        nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
        pad = 42
        correct_output = []
        for t in ts:
            next_output = torch.ones_like(ts[2]) * pad
            correct_output.append(next_output)
            next_output[: t.size(0)].copy_(t)
        correct_output = torch.stack(correct_output)
        padded = torch.nested.to_padded_tensor(nt, pad)
        self.assertEqual(padded, correct_output)

    @dtypes(torch.float, torch.float16, torch.double)
    def test_to_padded_tensor_dim3(self, device, dtype):
        ts = [
            torch.randn(16, 21, device=device, dtype=dtype),
            torch.randn(24, 32, device=device, dtype=dtype),
            torch.randn(40, 53, device=device, dtype=dtype),
        ]
        nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
        pad = 42
        correct_output = []
        for t in ts:
            next_output = torch.ones_like(ts[2]) * pad
            correct_output.append(next_output)
            next_output[: t.size(0), : t.size(1)].copy_(t)
        correct_output = torch.stack(correct_output)
        padded = torch.nested.to_padded_tensor(nt, pad)
        self.assertEqual(padded, correct_output)

    @dtypes(torch.float, torch.float16, torch.double)
    def test_to_padded_tensor_dim4(self, device, dtype):
        ts = [
            torch.randn(16, 21, 13, device=device, dtype=dtype),
            torch.randn(24, 32, 14, device=device, dtype=dtype),
            torch.randn(40, 53, 16, device=device, dtype=dtype),
        ]
        nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
        pad = 42
        correct_output = []
        for t in ts:
            next_output = torch.ones_like(ts[2]) * pad
            correct_output.append(next_output)
            next_output[: t.size(0), : t.size(1), : t.size(2)].copy_(t)
        correct_output = torch.stack(correct_output)
        padded = torch.nested.to_padded_tensor(nt, pad)
        self.assertEqual(padded, correct_output)

    # TODO: test noncontiguous to_padded_tensor
    # For now this tests the functionality of noncontiguous_to_padded_tensor
    # and the error message of to_padded_tensor
    # since to_padded_tensor does not support noncontiguous buffer yet
    @dtypes(torch.float, torch.float16, torch.double)
    @torch.inference_mode()
    def test_to_padded_tensor_noncontiguous(self, device, dtype):
        nt_contiguous, nt_noncontiguous = random_nt_noncontiguous_pair(
            (2, 3, 6, 7), device, dtype
        )
        # test noncontiguous_to_padded_tensor functionality
        self.assertEqual(
            torch.nested.to_padded_tensor(nt_contiguous, 0.0),
            noncontiguous_to_padded_tensor(nt_noncontiguous),
        )
        # test to_padded_tensor error message
        self.assertRaisesRegex(
            RuntimeError,
            r"for now to_padded_tensor only supports contiguous nested tensor",
            lambda: torch.nested.to_padded_tensor(nt_noncontiguous, 0.0),
        )

    @skipMeta
    def test_device_checks(self, device):
        nt = torch.nested.nested_tensor([], device=device)
        is_cuda = "cuda" in str(device)
        self.assertEqual(nt.is_cuda, is_cuda)

    @dtypes(torch.float, torch.float16, torch.double)
    def test_nested_tensor_indexing(self, device, dtype):
        # edge case: empty nested tensor
        nt0 = torch.nested.nested_tensor([])
        self.assertRaises(IndexError, lambda: nt0[0])
        # normal case
        x0 = torch.randn((2, 5), device=device, dtype=dtype)
        x1 = torch.randn((3, 4), device=device, dtype=dtype)
        nt = torch.nested.nested_tensor([x0, x1])
        # single index: only support integer in the batch dimension
        self.assertEqual(nt[0], x0)
        self.assertEqual(nt[-1], x1)
        self.assertRaises(IndexError, lambda: nt[2])
        self.assertRaises(IndexError, lambda: nt[-3])
        self.assertRaises(NotImplementedError, lambda: nt[:])
        self.assertEqual(nt[...], nt)
        # tuple of indices: only support integer in the batch dimension
        #                 + all possible indexing in the original tensor dimensions
        self.assertEqual(nt[0, 0, 0], x0[0, 0])
        self.assertEqual(nt[0, 1, :], x0[1, :])
        self.assertEqual(nt[1, ...], x1)
        self.assertRaises(IndexError, lambda: nt[1, 4, 2])
        self.assertRaises(NotImplementedError, lambda: nt[:, 1, 1])
        # test select on non-batch dimensions
        self.assertEqual(nt.select(1, 0)[0], x0.select(0, 0))
        self.assertEqual(nt.select(1, 0)[1], x1.select(0, 0))
        self.assertRaises(IndexError, lambda: nt.select(1, 3))
        self.assertEqual(nt.select(2, 0)[0], x0.select(1, 0))
        self.assertEqual(nt.select(2, 0)[1], x1.select(1, 0))
        self.assertRaises(IndexError, lambda: nt.select(2, 5))
        # make sure indexing returns a view
        nt[0].fill_(100.0)
        answer = torch.tensor(100.0, device=device, dtype=dtype).expand((2, 5))
        self.assertEqual(nt[0], answer)
        nt[1, 1, :].fill_(200.0)
        answer = torch.tensor(200.0, device=device, dtype=dtype).expand(4)
        self.assertEqual(nt[1, 1, :], answer)

        # Test that indexing works when requires_grad_(True)
        # previously this was failing because the backward kernel for select.int uses .sizes()
        nt = torch.nested.nested_tensor([x0, x1]).requires_grad_(True)
        self.assertEqual(nt[0], x0)
        self.assertEqual(nt[-1], x1)
        grad_x0 = torch.randn((2, 5), device=device, dtype=dtype)
        nt[0].backward(grad_x0)
        expected_grad = torch.nested.nested_tensor(
            [grad_x0, torch.zeros((3, 4), device=device, dtype=dtype)]
        )
        self.assertEqual(nt.grad, expected_grad)

    @parametrize(
        "func",
        [
            subtest(torch.nn.functional.relu, name="relu"),
            subtest(torch.nn.functional.relu_, name="relu_"),
            subtest(torch.nn.functional.gelu, name="gelu"),
            subtest(torch._C._nn.gelu_, name="gelu_"),
            subtest(torch.tanh, name="tanh"),
            subtest(torch.tanh_, name="tanh_"),
            subtest(torch.neg, name="neg"),
            subtest(torch.nn.functional.silu, name="silu"),
            subtest(partial(torch.nn.functional.silu, inplace=True), name="silu_"),
            subtest(torch.abs, name="abs"),
            subtest(torch.abs_, name="abs_"),
            subtest(torch.sgn, name="sgn"),
            subtest(torch.logical_not, name="logical_not"),
            subtest(torch.sin, name="sin"),
            subtest(torch.cos, name="cos"),
        ],
    )
    def test_activations(self, device, func):
        nt, nt_noncontiguous = random_nt_noncontiguous_pair(
            (2, 3, 6, 7), device=device, dtype=torch.float32
        )
        nested_result = func(nt)
        self.assertTrue(nested_result.is_nested)
        for t, t_res in zip(nt.unbind(), nested_result.unbind()):
            self.assertEqual(func(t), t_res)
        self.assertRaisesRegex(
            RuntimeError,
            "NestedTensor must be contiguous to get buffer.",
            lambda: func(nt_noncontiguous),
        )

    @parametrize("func", [subtest(torch.ge, name="ge"), subtest(torch.eq, name="eq")])
    def test_binary_ops_with_scalar(self, device, func):
        nt_contiguous, nt_noncontiguous = random_nt_noncontiguous_pair(
            (2, 3, 6, 7), device=device, dtype=torch.float32
        )
        scalar = 0.0

        # should work regardless of contiguity
        for nt in (nt_contiguous, nt_noncontiguous):
            nested_result = func(nt, scalar)
            self.assertTrue(nested_result.is_nested)
            for t, t_res in zip(nt.unbind(), nested_result.unbind()):
                self.assertEqual(func(t, scalar), t_res)

    @dtypes(*floating_types_and_half())
    def test_nested_tensor_chunk(self, device, dtype):
        # Transformer use case
        a = torch.randn(3, 3 * 4, device=device, dtype=dtype)
        b = torch.randn(2, 3 * 4, device=device, dtype=dtype)
        c = torch.randn(1, 3 * 4, device=device, dtype=dtype)
        a_chunks = a.chunk(3, dim=-1)
        b_chunks = b.chunk(3, dim=-1)
        c_chunks = c.chunk(3, dim=-1)

        a_nt = [a_chunks[0], b_chunks[0], c_chunks[0]]
        b_nt = [a_chunks[1], b_chunks[1], c_chunks[1]]
        c_nt = [a_chunks[2], b_chunks[2], c_chunks[2]]

        nt = torch.nested.nested_tensor([a, b, c])
        chunked = nt.chunk(3, dim=-1)

        self.assertEqual(chunked[0], torch.nested.nested_tensor(a_nt))
        self.assertEqual(chunked[1], torch.nested.nested_tensor(b_nt))
        self.assertEqual(chunked[2], torch.nested.nested_tensor(c_nt))

        for chunk in chunked:
            self.assertFalse(chunk.is_contiguous())

        # Failure chunking on ragged dimensions
        self.assertRaisesRegex(
            RuntimeError,
            "Chunk for nested tensors is currently only supported for the last dimension.",
            lambda: torch.chunk(nt, 5, dim=1),
        )
        self.assertRaisesRegex(
            RuntimeError,
            "Chunk for nested tensors is currently only supported for the last dimension.",
            lambda: torch.chunk(nt, 5, dim=0),
        )

        # Failure on non-contiguous nt
        _, nt_noncontiguous = random_nt_noncontiguous_pair((2, 3), device, dtype)
        self.assertRaisesRegex(
            RuntimeError,
            "chunk expects `self` to be contiguous.",
            lambda: torch.chunk(nt_noncontiguous, 5, dim=-1),
        )

        # Failure when calling non divisible n_chunks
        self.assertRaisesRegex(
            RuntimeError,
            "Chunk for nested tensors is only supported for "
            "nested tensors with trailing dimension divisible by chunks.",
            lambda: torch.chunk(nt, 5, dim=-1),
        )

        # Failure when calling backward on a chunk
        a = torch.randn(3, 3 * 4, device=device, dtype=dtype, requires_grad=True)
        b = torch.randn(2, 3 * 4, device=device, dtype=dtype, requires_grad=True)
        nt_grad = torch.nested.as_nested_tensor([a, b])
        chunked = torch.chunk(nt_grad, 2, dim=-1)
        self.assertRaisesRegex(
            RuntimeError,
            "Nested Strided Tensor doesn't support chunk backward.",
            lambda: chunked[0].backward(chunked[0].clone()),
        )

    @dtypes(*floating_types_and_half())
    def test_nested_tensor_split_with_sizes(self, device, dtype):
        a = torch.randn(3, 20, device=device, dtype=dtype)
        b = torch.randn(2, 20, device=device, dtype=dtype)
        c = torch.randn(1, 20, device=device, dtype=dtype)

        split_sizes = [4, 6, 10]
        a_splits = a.split_with_sizes(split_sizes, dim=-1)
        b_splits = b.split_with_sizes(split_sizes, dim=-1)
        c_splits = c.split_with_sizes(split_sizes, dim=-1)

        nt = torch.nested.nested_tensor([a, b, c])
        nt_splits = nt.split_with_sizes(split_sizes, dim=-1)

        for i, nt_split in enumerate(nt_splits):
            self.assertEqual(
                nt_split,
                torch.nested.nested_tensor([a_splits[i], b_splits[i], c_splits[i]]),
            )
            dense_strides = torch.stack(
                [
                    torch.tensor(a_splits[i].stride()),
                    torch.tensor(b_splits[i].stride()),
                    torch.tensor(c_splits[i].stride()),
                ]
            )
            self.assertEqual(nt_split._nested_tensor_strides(), dense_strides)
            self.assertFalse(nt_split.is_contiguous())

        # Failure calling on ragged dimensions
        self.assertRaisesRegex(
            RuntimeError,
            "split_with_sizes for nested tensors is currently only supported for the last dimension.",
            lambda: torch.split_with_sizes(nt, split_sizes, dim=1),
        )

        # Failure calling on non-last dimension
        self.assertRaisesRegex(
            RuntimeError,
            "split_with_sizes for nested tensors is currently only supported for the last dimension.",
            lambda: torch.split_with_sizes(nt, split_sizes, dim=0),
        )

        # Failure on non-contiguous nt
        _, nt_noncontiguous = random_nt_noncontiguous_pair((2, 3), device, dtype)
        self.assertRaisesRegex(
            RuntimeError,
            "split_with_sizes expects `self` to be contiguous.",
            lambda: torch.split_with_sizes(nt_noncontiguous, split_sizes, dim=-1),
        )

        # Failure when calling with split_sizes that don't cover the full dim size
        bad_split_sizes = [4, 6, 9]  # don't add up to 20
        self.assertRaisesRegex(
            RuntimeError,
            "split_with_sizes expects split_sizes to sum exactly to 20",
            lambda: torch.split_with_sizes(nt, bad_split_sizes, dim=-1),
        )

    @dtypes(torch.float, torch.float16, torch.double)
    @torch.inference_mode()
    def test_nested_tensor_indexing_noncontiguous(self, device, dtype):
        nt_contiguous, nt_noncontiguous = random_nt_noncontiguous_pair(
            (2, 3, 6, 7), device, dtype
        )
        self.assertEqual(nt_contiguous.size(0), nt_noncontiguous.size(0))
        n = nt_contiguous.size(0)
        for i in range(n):
            self.assertEqual(nt_contiguous[i], nt_noncontiguous[i])

    @dtypes(torch.float, torch.float16)
    @skipMeta
    @torch.inference_mode()
    @parametrize("transpose", [True, False])
    def test_nested_tensor_add(self, device, dtype, transpose):
        if transpose:
            a = torch.randn(2, 2, 2, device=device, dtype=dtype)
            b = torch.rand(2, 2, 2, device=device, dtype=dtype)
            c = a.transpose(-1, -2).contiguous()
            d = b.transpose(-1, -2).contiguous()
            nt1 = torch.nested.nested_tensor([a, b, a, b])
            nt2 = torch.nested.nested_tensor([c, d, c, d]).transpose(-1, -2)
        else:
            (nt1, nt2) = self.random_nt_pair(device, dtype, 4, (4, 4))
        ref = torch.nested.nested_tensor(
            [t1 + t2 for (t1, t2) in zip(nt1.unbind(), nt2.unbind())]
        )
        out = nt1 + nt2
        self.assertEqual(ref, out)

    @dtypes(torch.float, torch.float16)
    @skipMeta
    @torch.inference_mode()
    @parametrize("transpose", [True, False])
    def test_nested_tensor_sub(self, device, dtype, transpose):
        if transpose:
            a = torch.randn(2, 2, 2, device=device, dtype=dtype)
            b = torch.rand(2, 2, 2, device=device, dtype=dtype)
            c = a.transpose(-1, -2).contiguous()
            d = b.transpose(-1, -2).contiguous()
            nt1 = torch.nested.nested_tensor([a, b, a, b])
            nt2 = torch.nested.nested_tensor([c, d, c, d]).transpose(-1, -2)
        else:
            (nt1, nt2) = self.random_nt_pair(device, dtype, 4, (4, 4))
        ref = torch.nested.nested_tensor(
            [t1 - t2 for (t1, t2) in zip(nt1.unbind(), nt2.unbind())]
        )
        out = nt1 - nt2
        self.assertEqual(ref, out)

    @onlyCUDA
    @dtypes(torch.float, torch.float16)
    @torch.inference_mode()
    @parametrize("embedding_dim", [8, 128, 256, 384])
    def test_nested_tensor_dense_elementwise(self, device, dtype, embedding_dim):
        def _test_add_mul(nt, t):
            ref_add = torch.nested.nested_tensor(
                [t1 + t2 for (t1, t2) in zip(nt.unbind(), t.unbind())]
            )
            ref_mul = torch.nested.nested_tensor(
                [t1 * t2 for (t1, t2) in zip(nt.unbind(), t.unbind())]
            )
            self.assertEqual(nt.add(t), ref_add)
            self.assertEqual(nt.mul(t), ref_mul)

        batch_size = 32
        seq_lens = torch.randint(low=0, high=10, size=(batch_size,))

        # [B, *, D], [B, 1, D] case
        ts = [torch.randn((seq_len, embedding_dim)) for seq_len in seq_lens]
        nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
        t = torch.randn((batch_size, 1, embedding_dim), device=device, dtype=dtype)
        _test_add_mul(nt, t)

        # [B, *], [B, 1] case
        ts = [torch.randn(seq_len) for seq_len in seq_lens]
        nt = torch.nested.nested_tensor(ts, device=device, dtype=dtype)
        t = torch.randn((batch_size, 1), device=device, dtype=dtype)
        _test_add_mul(nt, t)

    @dtypes(torch.float, torch.float16)
    @skipMeta
    @torch.inference_mode()
    def test_nested_tensor_mul(self, device, dtype):
        # nested tensor * nested tensor
        (nt1, nt2) = self.random_nt_pair(device, dtype, 4, (4, 4))
        ref = torch.nested.nested_tensor(
            [t1 * t2 for (t1, t2) in zip(nt1.unbind(), nt2.unbind())]
        )
        out = nt1 * nt2
        self.assertEqual(ref, out)
        # nested tensor * scalar
        number = 10.0
        scalar = torch.tensor(number).to(dtype).to(device)
        ref = torch.nested.nested_tensor([t * number for t in nt1.unbind()])
        out_number0 = nt1 * number
        out_number1 = number * nt1
        out_scalar0 = nt1 * scalar
        out_scalar1 = scalar * nt1
        self.assertEqual(out_number0, ref)
        self.assertEqual(out_number1, ref)
        self.assertEqual(out_scalar0, ref)
        self.assertEqual(out_scalar1, ref)
        # error case: numel == 1 but dim > 0
        vector = torch.tensor([number]).to(dtype).to(device)
        self.assertRaisesRegex(
            RuntimeError,
            "Expected both self and other to be nested, but got a nested self and non-nested other",
            lambda: nt1.mul(vector),
        )
        self.assertRaisesRegex(
            RuntimeError,
            "Expected both self and other to be nested, but got a non-nested self and nested other",
            lambda: vector.mul(nt1),
        )

    @dtypes(torch.float, torch.float16)
    @skipMeta
    @torch.inference_mode()
    def test_nested_tensor_div(self, device, dtype):
        nt, nt2 = self.random_nt_pair(device, dtype, 4, (4, 4))
        scale = 4.0
        ref = torch.nested.nested_tensor([t / scale for t in nt.unbind()])
        out = nt / 4.0
        self.assertEqual(ref, out)
        ref_transposed = ref.transpose(1, 2)
        out = nt.transpose(1, 2) / 4.0
        self.assertEqual(ref_transposed, out)

        ref = torch.nested.nested_tensor(
            [t / t2 for (t, t2) in zip(nt.unbind(), nt2.unbind())]
        )
        out = nt / nt2
        self.assertEqual(ref, out)

        out = nt.transpose(1, 2) / nt2.transpose(1, 2)
        self.assertEqual(ref.transpose(1, 2), out)

        nt_transpose_copy = torch.nested.nested_tensor(
            [t.transpose(0, 1) for t in nt.unbind()]
        )

        self.assertRaisesRegex(
            RuntimeError,
            "div requires strides to match when given NestedTensors",
            lambda: nt_transpose_copy.transpose(1, 2) / nt2,
        )

        nt = torch.nested.nested_tensor(
            [torch.randn(i, 4) for i in [3, 4, 5]], device=device, dtype=dtype
        )
        nt_chunks = nt.chunk(2, -1)
        self.assertRaisesRegex(
            RuntimeError,
            "div requires offsets to match when given NestedTensors",
            lambda: nt_chunks[0] / nt_chunks[1],
        )

    @dtypes(torch.float, torch.float16)
    @skipMeta
    @torch.inference_mode()
    def test_nested_tensor_add_in_place(self, device, dtype):
        (nt1, nt2) = self.random_nt_pair(device, dtype, 4, (4, 4))
        ref = torch.nested.nested_tensor(
            [t1 + t2 for (t1, t2) in zip(nt1.unbind(), nt2.unbind())]
        )
        nt1 += nt2
        self.assertEqual(ref, nt1)

    @dtypes(torch.float, torch.float16)
    @skipMeta
    @torch.inference_mode()
    def test_nested_tensor_mul_in_place(self, device, dtype):
        # nested tensor * nested tensor
        (nt1, nt2) = self.random_nt_pair(device, dtype, 4, (4, 4))
        ref = torch.nested.nested_tensor(
            [t1 * t2 for (t1, t2) in zip(nt1.unbind(), nt2.unbind())]
        )
        nt1 *= nt2
        self.assertEqual(ref, nt1)
        # nested tensor * scalar
        number = 10.0
        scalar = torch.tensor(number).to(dtype).to(device)
        ref = torch.nested.nested_tensor([t * number for t in nt1.unbind()])
        out_number = nt1.clone()
        out_number *= number
        out_scalar = nt1.clone()
        out_scalar *= scalar
        self.assertEqual(out_number, ref)
        self.assertEqual(out_scalar, ref)
        self.assertRaisesRegex(
            RuntimeError,
            r"output with shape \[.*\] doesn't match the broadcast shape \[.*\]",
            lambda: scalar.mul_(nt1),
        )
        # error case: numel == 1 but dim > 0
        vector = torch.tensor([number]).to(dtype).to(device)
        self.assertRaisesRegex(
            RuntimeError,
            "Expected both self and other to be nested, but got a nested self and non-nested other",
            lambda: nt1.mul_(vector),
        )
        self.assertRaisesRegex(
            RuntimeError,
            "Expected both self and other to be nested, but got a non-nested self and nested other",
            lambda: vector.mul_(nt1),
        )

    @onlyCPU
    @skipMeta
    @dtypes(torch.float)
    def test_nested_tensor_sum_dim(self, device, dtype):
        params = ((2, (1, 1)), ((4), (4, 4)), (10, (3, 5, 7)))

        def test_sum(device, dtype, ntensors, max_sizes, dim, keepdim=True):
            nt = random_nt(device, dtype, ntensors, max_sizes, require_non_empty=False)
            nt2 = nt.clone()
            ub2 = nt2.unbind()
            nt.requires_grad_(True)
            [t.requires_grad_(True) for t in ub2]
            nt_sum = nt.sum(dim=dim, keepdim=keepdim)
            ub2_sum = [t.sum(-1, keepdim=keepdim) for t in ub2]
            self.assertEqual(nt_sum, torch.nested.nested_tensor(ub2_sum))

            # test backward
            # generate gradient tensor that has the same size as the output
            size = nt_sum._nested_tensor_size()
            gt2 = []
            for i in range(ntensors):
                gt2.append(torch.randn(size[i].tolist(), device=device, dtype=dtype))
            gt = torch.nested.nested_tensor(gt2).clone()
            nt_sum.backward(gt)
            for t2, g2 in zip(ub2_sum, gt2):
                t2.backward(g2)
            self.assertEqual(nt.grad, torch.nested.nested_tensor([t.grad for t in ub2]))
            return

        for ntensors, max_sizes in params:
            test_sum(device, dtype, ntensors, max_sizes, len(max_sizes))

        # Test error inputs
        with self.assertRaisesRegex(
            RuntimeError, "NestedTensor can only be reduced across the last"
        ):
            torch.nested.nested_tensor(
                [torch.tensor([3, 4, 5]), torch.tensor([1, 2])]
            ).sum(0, keepdim=True)

        with self.assertRaisesRegex(
            RuntimeError, "NestedTensor only allows reduction of a single"
        ):
            torch.nested.nested_tensor(
                [torch.tensor([[3, 4, 5]]), torch.tensor([[1, 2]])]
            ).sum([0, 1], keepdim=True)

        with self.assertRaisesRegex(
            RuntimeError, "NestedTensor always requires keepdim=True for now."
        ):
            torch.nested.nested_tensor(
                [torch.tensor([3, 4, 5]), torch.tensor([1, 2])]
            ).sum(-1)

    @dtypes(torch.float, torch.float16)
    def test_contiguous(self, device, dtype):
        # Since we don't have access to the buffer in python this is harder to show what
        # we are testing for. When we call chunk on a consistent dim of a NT
        # for chunk_size > 1 the resulting tensors are views of the original NT
        # whose numels is now less than the size of the buffer. Clone was
        # previously creating a new NT with a buffer that was the same size as the
        # original.
        nt_contiguous = torch.nested.nested_tensor(
            [
                torch.randn(2, 20, device=device, dtype=dtype),
                torch.randn(4, 20, device=device, dtype=dtype),
            ]
        )
        # Split up the last dimension which has a consistent size of 20 into 5 chunks
        chunks = nt_contiguous.chunk(5, dim=-1)

        # # Check chunks are contiguous after calling contiguous
        for chunk in chunks:
            self.assertFalse(chunk.is_contiguous())
            self.assertTrue(chunk.contiguous().is_contiguous())

    @dtypes(torch.float, torch.float16)
    @skipMeta
    def test_clone(self, device, dtype):
        nt1 = random_nt(device, dtype, 4, (4, 4), (1, 1))
        nt2 = nt1.clone()
        # Verify the values match
        self.assertEqual(nt1, nt2)
        # Verify modifying nt2 doesn't affect nt1
        nt2.mul_(nt1)
        ub1 = nt1.unbind()
        ub2 = nt2.unbind()
        for i in range(len(ub1)):
            self.assertNotEqual(ub1[i], ub2[i])

        nt1.clone(memory_format=torch.preserve_format)
        msg = "Nested tensor clone supports Preserve and Contiguous memory formats, called clone with memory format: ChannelsLast"
        with self.assertRaisesRegex(RuntimeError, msg):
            nt1.clone(memory_format=torch.channels_last)

    # cannot test torch.float16 because: RuntimeError: "bernoulli_scalar_cpu_" not implemented for 'Half'
    @decorateIf(xfailIfTorchDynamo, lambda params: params["layout"] == torch.jagged)
    @dtypes(torch.float, torch.double)
    @parametrize("layout", [torch.strided, torch.jagged], name_fn=layout_name)
    def test_dropout(self, device, dtype, layout):
        # edge case: empty nested tensor
        # TODO: support empty NT in jagged layout
        if layout == torch.strided:
            nt0 = torch.nested.nested_tensor([], layout=layout)
            y = torch.nn.functional.dropout(nt0, 0.5)
            self.assertEqual(nt0, y)
        # normal nested tensor
        ntensors = 4
        if layout == torch.jagged:
            nt = random_nt(device, dtype, ntensors, (4, 4), (0, 3), layout=layout)
        else:
            nt = random_nt(device, dtype, ntensors, (4, 4), layout=layout)
        # edge case: invalid dropout
        self.assertRaises(ValueError, lambda: torch.nn.Dropout(-0.1))
        self.assertRaises(ValueError, lambda: torch.nn.Dropout(1.1))
        self.assertRaises(ValueError, lambda: torch.nn.functional.dropout(nt, -0.1))
        self.assertRaises(ValueError, lambda: torch.nn.functional.dropout(nt, 1.1))
        # edge case: no dropout
        dropouter = torch.nn.Dropout(0.0)
        y0 = dropouter(nt)
        y1 = torch.nn.functional.dropout(nt, 0.0)
        self.assertEqual(nt, y0)
        self.assertEqual(nt, y1)
        # edge case: all dropout
        dropouter = torch.nn.Dropout(1.0)
        y0 = dropouter(nt)
        y1 = torch.nn.functional.dropout(nt, 1.0)
        nt0 = torch.zeros_like(nt)
        self.assertEqual(nt0, y0)
        self.assertEqual(nt0, y1)
        # normal case: normal dropout
        p = 0.2
        y = torch.nn.functional.dropout(nt, p)
        expect = nt.clone()
        if layout == torch.jagged:
            expect = torch.where(y == 0.0, y, nt)
            expect /= 1.0 - p
            self.assertEqual(y, expect)
        else:
            expect = nt.clone()
            for i in range(ntensors):
                actual_tensor = y[i].view(-1)
                expect_tensor = expect[i].view(-1)
                for j in range(actual_tensor.shape[0]):
                    if actual_tensor[j].item() == 0.0:
                        expect_tensor[j] = 0.0
                    else:
                        expect_tensor[j] /= 1.0 - p
            self.assertEqual(y, expect)
        with freeze_rng_state():
            dropouter = torch.nn.Dropout(p)
            y0 = dropouter(nt)
        with freeze_rng_state():
            y1 = torch.nn.functional.dropout(nt, p)
        self.assertEqual(y0, y1)

    @dtypes(torch.float, torch.double)
    def test_dropout_noncontiguous(self, device, dtype):
        ntensors = 4
        nt0 = random_nt(device, dtype, ntensors, (4, 4))
        nt1 = nt0.transpose(-1, -2)
        p = 0.3
        with freeze_rng_state():
            dropouter = torch.nn.Dropout(p)
            y0 = dropouter(nt0)
        with freeze_rng_state():
            y1 = torch.nn.functional.dropout(nt1, p).transpose(-1, -2)
        self.assertEqual(y0, y1)

    # cannot test torch.float16 because: RuntimeError: "softmax_kernel_impl" not implemented for 'Half'
    @dtypes(torch.float, torch.double)
    def test_softmax(self, device, dtype):
        # normal nested tensor
        ntensors = 4
        nt = random_nt(device, dtype, ntensors, (4, 4))
        # error case: softmax across nested dimension
        self.assertRaisesRegex(
            RuntimeError,
            "Cannot apply softmax across nested dimension 0",
            lambda: torch.nn.functional.softmax(nt, 0),
        )
        self.assertRaisesRegex(
            RuntimeError,
            "Cannot apply softmax across nested dimension 0",
            lambda: torch.nn.functional.softmax(nt, -3),
        )
        # error case: dimension out of range
        self.assertRaises(IndexError, lambda: torch.nn.functional.softmax(nt, 3))
        self.assertRaises(IndexError, lambda: torch.nn.functional.softmax(nt, -4))
        # normal case: should equal to padding -inf
        softmaxer = torch.nn.Softmax(1)
        y0 = softmaxer(nt)
        y1 = torch.nn.functional.softmax(nt, 1)
        self.assertEqual(y0, y1)
        pt = torch.nested.to_padded_tensor(nt, float("-inf"))
        # if an entire slice is padded, then softmax will return 0.0 / 0.0 = nan
        # however, physically speaking that should be 0.0
        expect = torch.nn.functional.softmax(pt, 1).nan_to_num_(0.0)
        self.assertEqual(torch.nested.to_padded_tensor(y0, 0.0), expect)
        # edge case: empty nested tensor
        nt0 = torch.nested.nested_tensor([])
        y = torch.nn.functional.softmax(nt0, 1)
        self.assertEqual(nt0, y)
        # edge case: nesting scalars
        nt1 = torch.nested.nested_tensor([torch.tensor(0.0), torch.tensor(1.0)])
        self.assertRaises(RuntimeError, lambda: torch.nn.functional.softmax(nt1, 0))
        self.assertRaises(IndexError, lambda: torch.nn.functional.softmax(nt1, 1))

    @dtypes(torch.float, torch.double)
    @torch.inference_mode()
    def test_softmax_noncontiguous(self, device, dtype):
        nt_contiguous, nt_noncontiguous = random_nt_noncontiguous_pair(
            (2, 3, 6, 7), device, dtype
        )
        self.assertEqual(
            torch.nn.functional.softmax(nt_contiguous, -1),
            torch.nn.functional.softmax(nt_noncontiguous, -1),
        )

    def _test_bmm(self, device, dtype):
        # error case: not 3D tensors
        nt0 = torch.nested.nested_tensor([], device=device, dtype=dtype)
        nt1 = torch.nested.nested_tensor(
            [torch.randn(2), torch.randn(3)], device=device, dtype=dtype
        )
        nt2 = torch.nested.nested_tensor(
            [torch.randn((2, 4)), torch.randn((3, 4))], device=device, dtype=dtype
        )
        self.assertRaisesRegex(
            RuntimeError, "batch1 must be a 3D tensor", lambda: nt0.bmm(nt0)
        )
        self.assertRaisesRegex(
            RuntimeError, "batch1 must be a 3D tensor", lambda: nt0.bmm(nt1)
        )
        self.assertRaisesRegex(
            RuntimeError, "batch1 must be a 3D tensor", lambda: nt0.bmm(nt2)
        )
        self.assertRaisesRegex(
            RuntimeError, "batch1 must be a 3D tensor", lambda: nt1.bmm(nt0)
        )
        self.assertRaisesRegex(
            RuntimeError, "batch1 must be a 3D tensor", lambda: nt1.bmm(nt1)
        )
        self.assertRaisesRegex(
            RuntimeError, "batch1 must be a 3D tensor", lambda: nt1.bmm(nt2)
        )
        self.assertRaisesRegex(
            RuntimeError, "batch2 must be a 3D tensor", lambda: nt2.bmm(nt0)
        )
        self.assertRaisesRegex(
            RuntimeError, "batch2 must be a 3D tensor", lambda: nt2.bmm(nt1)
        )
        # error case: incompatible batch size
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 4)), torch.randn((3, 4))], device=device, dtype=dtype
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((4, 6)), torch.randn((4, 5)), torch.randn((4, 7))],
            device=device,
            dtype=dtype,
        )
        self.assertRaisesRegex(
            RuntimeError,
            "Expected size for the 1st dimension of batch2 tensor to be: 2 but got: 3.",
            lambda: nt0.bmm(nt1),
        )
        self.assertRaisesRegex(
            RuntimeError,
            "Expected size for the 1st dimension of batch2 tensor to be: 3 but got: 2.",
            lambda: nt1.bmm(nt0),
        )
        # error case: underlying matrices cannot be multiplied
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 4)), torch.randn((3, 4))], device=device, dtype=dtype
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"0-th nested matrices in batch cannot be multiplied \(2x4 and 2x4\)",
            lambda: nt0.bmm(nt0),
        )
        # normal nested tensor
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 4)), torch.randn((3, 7))], device=device, dtype=dtype
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((4, 6)), torch.randn((7, 5))], device=device, dtype=dtype
        )
        actual = torch.nested.to_padded_tensor(nt0.bmm(nt1), 0.0)
        expect = torch.nested.to_padded_tensor(nt0, 0.0).bmm(
            torch.nested.to_padded_tensor(nt1, 0.0)
        )
        if dtype == torch.float16:
            self.assertEqual(actual, expect, rtol=1e-3, atol=1e-3)
        else:
            self.assertEqual(actual, expect)

        # nested tensor bmm normal tensor
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 7)), torch.randn((3, 7))], device=device, dtype=dtype
        )
        nt1 = torch.rand(2, 7, 5, dtype=dtype, device=device)
        actual = torch.nested.to_padded_tensor(nt0.bmm(nt1), 0.0)
        expect = torch.nested.to_padded_tensor(nt0, 0.0).bmm(nt1)
        if dtype == torch.float16:
            self.assertEqual(actual, expect, rtol=1e-3, atol=1e-3)
        else:
            self.assertEqual(actual, expect)

        # nested tensor bmm normal tensor with non-contiguous view
        nt1 = torch.rand(2, 5, 7, dtype=dtype, device=device)
        nt1 = nt1.transpose(1, 2)
        actual = torch.nested.to_padded_tensor(nt0.bmm(nt1), 0.0)
        expect = torch.nested.to_padded_tensor(nt0, 0.0).bmm(nt1)
        if dtype == torch.float16:
            self.assertEqual(actual, expect, rtol=1e-3, atol=1e-3)
        else:
            self.assertEqual(actual, expect)

        # normal tensor bmm nested tensor
        nt0 = torch.rand(2, 5, 7, dtype=dtype, device=device)
        nt1 = torch.nested.nested_tensor(
            [torch.randn((7, 6)), torch.randn((7, 5))], device=device, dtype=dtype
        )
        actual = torch.nested.to_padded_tensor(nt0.bmm(nt1), 0.0)
        expect = nt0.bmm(torch.nested.to_padded_tensor(nt1, 0.0))
        if dtype == torch.float16:
            self.assertEqual(actual, expect, rtol=1e-3, atol=1e-3)
        else:
            self.assertEqual(actual, expect)

        # test tensorcore path
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 8)), torch.randn((3, 16))], device=device, dtype=dtype
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((8, 8)), torch.randn((16, 8))], device=device, dtype=dtype
        )
        actual = torch.nested.to_padded_tensor(nt0.bmm(nt1), 0.0)
        expect = torch.nested.to_padded_tensor(nt0, 0.0).bmm(
            torch.nested.to_padded_tensor(nt1, 0.0)
        )
        if dtype == torch.float16:
            self.assertEqual(actual, expect, rtol=1e-3, atol=1e-3)
        else:
            self.assertEqual(actual, expect)

    @onlyCUDA
    @dtypes(torch.float, torch.double, torch.float16)
    def test_bmm_cuda(self, device, dtype):
        self._test_bmm(device, dtype)

    @onlyCPU
    # cannot test torch.float16 because: RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
    @dtypes(torch.float, torch.double)
    def test_bmm_cpu(self, device, dtype):
        self._test_bmm(device, dtype)

    # cannot test torch.float16 because: RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
    @dtypes(torch.float, torch.double)
    def test_bmm_noncontiguous(self, device, dtype):
        nt0_contiguous, nt0_noncontiguous = random_nt_noncontiguous_pair(
            (2, 3), device, dtype
        )
        nt1_contiguous, nt1_noncontiguous = random_nt_noncontiguous_pair(
            (6, 7), device, dtype
        )
        self.assertEqual(
            nt0_contiguous.transpose(-1, -2).bmm(nt1_contiguous),
            nt0_noncontiguous.transpose(-1, -2).bmm(nt1_noncontiguous),
        )

    @dtypes(torch.float, torch.double)
    def test_matmul_with_bmm_path(self, device, dtype):
        def unbind_rebind_matmul(nt1, nt2):
            t1s = nt1.unbind()
            t2s = nt2.unbind()
            out_ts = [t1.matmul(t2) for t1, t2 in zip(t1s, t2s)]
            return torch.nested.nested_tensor(out_ts)

        # [N, n_head, *, head_dim], [N, n_head, head_dim, *]
        Ns = [1, 2, 5]
        n_heads = np.random.randint(2, 5)
        head_dim = 3
        t1s = []
        t2s = []
        for N in Ns:
            for _ in range(N):
                seq_len1 = np.random.randint(2, 5)
                seq_len2 = np.random.randint(2, 5)
                t1s.append(torch.randn(n_heads, seq_len1, head_dim))
                t2s.append(torch.randn(n_heads, head_dim, seq_len2))
            nt1 = torch.nested.nested_tensor(t1s, device=device, dtype=dtype)
            nt2 = torch.nested.nested_tensor(t2s, device=device, dtype=dtype)
            self.assertEqual(torch.matmul(nt1, nt2), unbind_rebind_matmul(nt1, nt2))

        # test with noncontiguous
        t3s = []
        t4s = []
        for _ in range(N):
            seq_len = np.random.randint(2, 5)
            t3s.append(torch.randn(seq_len, n_heads, head_dim))
            t4s.append(torch.randn(seq_len, n_heads, head_dim))
        nt3 = torch.nested.nested_tensor(t3s, device=device, dtype=dtype).transpose(
            1, 2
        )
        nt4 = (
            torch.nested.nested_tensor(t4s, device=device, dtype=dtype)
            .transpose(1, 2)
            .transpose(2, 3)
        )
        self.assertEqual(torch.matmul(nt3, nt4), unbind_rebind_matmul(nt3, nt4))

    # cannot test torch.float16 because: RuntimeError: "bmm" not implemented for 'Half'
    @dtypes(torch.float, torch.double)
    def test_matmul(self, device, dtype):
        # error case: one is nested but the other is not
        nt = torch.nested.nested_tensor(
            [torch.randn(2), torch.randn(3)], device=device, dtype=dtype
        )
        t = torch.randn(4, device=device, dtype=dtype)
        self.assertRaisesRegex(
            RuntimeError,
            "Expected both to be nested, but got a nested self and non-nested other",
            lambda: torch.matmul(nt, t),
        )
        self.assertRaisesRegex(
            RuntimeError,
            "Expected both to be nested, but got a non-nested self and nested other",
            lambda: torch.matmul(t, nt),
        )
        # error case: not 3+D tensors
        nt0 = torch.nested.nested_tensor([], device=device, dtype=dtype)
        nt1 = torch.nested.nested_tensor(
            [torch.randn(2), torch.randn(3)], device=device, dtype=dtype
        )
        nt2 = torch.nested.nested_tensor(
            [torch.randn((2, 4)), torch.randn((3, 4))], device=device, dtype=dtype
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"matmul: For nested tensors, only inputs with >= 3 dims are currently supported. 1st input has rank: [0-9]+",
            lambda: torch.matmul(nt0, nt0),
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"matmul: For nested tensors, only inputs with >= 3 dims are currently supported. 1st input has rank: [0-9]+",
            lambda: torch.matmul(nt0, nt1),
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"matmul: For nested tensors, only inputs with >= 3 dims are currently supported. 1st input has rank: [0-9]+",
            lambda: torch.matmul(nt0, nt2),
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"matmul: For nested tensors, only inputs with >= 3 dims are currently supported. 1st input has rank: [0-9]+",
            lambda: torch.matmul(nt1, nt0),
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"matmul: For nested tensors, only inputs with >= 3 dims are currently supported. 1st input has rank: [0-9]+",
            lambda: torch.matmul(nt1, nt1),
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"matmul: For nested tensors, only inputs with >= 3 dims are currently supported. 1st input has rank: [0-9]+",
            lambda: torch.matmul(nt1, nt2),
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"matmul: For nested tensors, only inputs with >= 3 dims are currently supported. 2nd input has rank: [0-9]+",
            lambda: torch.matmul(nt2, nt0),
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"matmul: For nested tensors, only inputs with >= 3 dims are currently supported. 2nd input has rank: [0-9]+",
            lambda: torch.matmul(nt2, nt1),
        )
        # error case: incompatible batch size
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 4)), torch.randn((3, 4))], device=device, dtype=dtype
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((4, 6)), torch.randn((4, 5)), torch.randn((4, 7))],
            device=device,
            dtype=dtype,
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"matmul: Expected size for the 1st dimension of 2nd input tensor to be: [0-9]+ but got: [0-9]+.",
            lambda: torch.matmul(nt0, nt1),
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"matmul: Expected size for the 1st dimension of 2nd input tensor to be: [0-9]+ but got: [0-9]+.",
            lambda: torch.matmul(nt1, nt0),
        )
        # error case: incompatible (wrong) batch sizes that shouldn't even broadcast?
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 2, 4)), torch.randn((2, 3, 4))], device=device, dtype=dtype
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((3, 4, 6)), torch.randn((3, 4, 5))], device=device, dtype=dtype
        )
        self.assertRaisesRegex(
            RuntimeError,
            "matmul(): For nested tensors, batch dimensions must have the same sizes,",
            lambda: torch.matmul(nt0, nt1),
        )
        # error case: incompatible batch sizes that should technically broadcast
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 2, 4)), torch.randn((1, 3, 4))], device=device, dtype=dtype
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((1, 4, 6)), torch.randn((3, 4, 5))], device=device, dtype=dtype
        )
        self.assertRaisesRegex(
            RuntimeError,
            "matmul(): For nested tensors, batch dimensions must have the same sizes,",
            lambda: torch.matmul(nt0, nt1),
        )
        # error case: underlying matrices cannot be multiplied
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 4)), torch.randn((3, 4))], device=device, dtype=dtype
        )
        self.assertRaisesRegex(
            RuntimeError,
            "matmul(): Nested tensors cannot be matrix multiplied",
            lambda: torch.matmul(nt0, nt0),
        )
        # normal nested tensor: 3D
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 4)), torch.randn((3, 7))], device=device, dtype=dtype
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((4, 6)), torch.randn((7, 5))], device=device, dtype=dtype
        )
        actual = torch.nested.to_padded_tensor(torch.matmul(nt0, nt1), 0.0)
        expect = torch.matmul(
            torch.nested.to_padded_tensor(nt0, 0.0),
            torch.nested.to_padded_tensor(nt1, 0.0),
        )
        self.assertEqual(actual, expect)
        # normal nested tensor: 4D (with testing for batch_size=1)
        nt0 = torch.nested.nested_tensor(
            [torch.randn((1, 2, 4)), torch.randn((8, 3, 7))], device=device, dtype=dtype
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((1, 4, 6)), torch.randn((8, 7, 5))], device=device, dtype=dtype
        )
        actual = torch.nested.to_padded_tensor(torch.matmul(nt0, nt1), 0.0)
        expect = torch.matmul(
            torch.nested.to_padded_tensor(nt0, 0.0),
            torch.nested.to_padded_tensor(nt1, 0.0),
        )
        self.assertEqual(actual, expect)
        # normal nested tensor: 5D
        nt0 = torch.nested.nested_tensor(
            [torch.randn((8, 9, 2, 4)), torch.randn((8, 9, 3, 7))],
            device=device,
            dtype=dtype,
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((8, 9, 4, 6)), torch.randn((8, 9, 7, 5))],
            device=device,
            dtype=dtype,
        )
        actual = torch.nested.to_padded_tensor(torch.matmul(nt0, nt1), 0.0)
        expect = torch.matmul(
            torch.nested.to_padded_tensor(nt0, 0.0),
            torch.nested.to_padded_tensor(nt1, 0.0),
        )
        self.assertEqual(actual, expect)

    # only supported on CUDA for now
    @dtypes(torch.float, torch.double)
    def test_matmul_nt_with_broadcasted_t(self, device, dtype):
        # NT (B, *, C, D) with T (D, E) broadcasting case
        nt = random_nt_from_dims([3, None, 4, 5], device=device, dtype=dtype)
        t = torch.randn(5, 6, device=device, dtype=dtype)
        output = torch.matmul(nt, t)

        # should be equivalent to matmul-ing each component with the dense tensor
        self.assertEqual(nt.size(0), output.size(0))
        for component, out_component in zip(nt, output):
            self.assertEqual(out_component, torch.matmul(component, t))

    # cannot test torch.float16 because: RuntimeError: "bmm" not implemented for 'Half'
    @dtypes(torch.float, torch.double)
    def test_matmul_noncontiguous(self, device, dtype):
        nt0_contiguous, nt0_noncontiguous = random_nt_noncontiguous_pair(
            (2, 3), device, dtype
        )
        nt1_contiguous, nt1_noncontiguous = random_nt_noncontiguous_pair(
            (6, 7), device, dtype
        )
        self.assertEqual(
            torch.matmul(nt0_contiguous.transpose(-1, -2), nt1_contiguous),
            torch.matmul(nt0_noncontiguous.transpose(-1, -2), nt1_noncontiguous),
        )

    @dtypes(torch.float, torch.double)
    def test_linear(self, device, dtype):
        a = torch.randn(1, 2, device=device, dtype=dtype)
        b = torch.randn(2, 2, device=device, dtype=dtype)
        c = torch.randn(3, 2, device=device, dtype=dtype)
        nt = torch.nested.nested_tensor([a, b, c])

        weight = torch.randn(2, 2, device=device, dtype=dtype)
        bias = torch.randn(2, device=device, dtype=dtype)
        # success case
        torch.functional.F.linear(nt, weight, bias)

        # invalid nested tensor dimension
        msg = r"Linear requires nested_tensor.dim == 3 and dense_matrix.dim == 2. Nested tensor dim: 2. Dense tensor dim: 2"
        nt1 = torch.nested.nested_tensor(
            [
                torch.randn(1, device=device, dtype=dtype),
                torch.randn(2, device=device, dtype=dtype),
            ]
        )
        with self.assertRaisesRegex(RuntimeError, msg):
            torch.functional.F.linear(nt1, weight, bias)

        # invalid weight shape
        msg = r"Linear requires nested_tensor.dim == 3 and dense_matrix.dim == 2. Nested tensor dim: 3. Dense tensor dim: 3"
        weight1 = torch.randn(2, 2, 3, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, msg):
            torch.functional.F.linear(nt, weight1, bias)

        # inconsistent last dim of nested tensor
        msg = r"Expected all tensors in nested tensor to have the same trailing dimension, instead last dimension equals:"
        nt2 = torch.nested.nested_tensor(
            [
                torch.randn(1, 2, device=device, dtype=dtype),
                torch.randn(2, 3, device=device, dtype=dtype),
            ]
        )
        with self.assertRaisesRegex(RuntimeError, msg):
            torch.functional.F.linear(nt2, weight, bias)

        # Mismatch of nested tensor last dim and weight dimension
        weight2 = torch.randn(2, 4, device=device, dtype=dtype)
        msg = (
            r"Shape mismatch for NestedTensor Linear: Expected input's \(a nested tensor\) 'last_dim'"
            r" to equal 'weight.size\(1\), but got: last_dim = 2, and weight.size\(1\) = 4"
        )
        with self.assertRaisesRegex(RuntimeError, msg):
            torch.functional.F.linear(nt, weight2, bias)

        # Nested tensor input and nested weight
        nt_weight = nt.clone()
        msg = r"Linear does not support nested weight when input is a nested tensor."
        with self.assertRaisesRegex(RuntimeError, msg):
            torch.functional.F.linear(nt, nt_weight, bias)

    # TODO: test noncontiguous linear
    # For now this tests the error message of linear
    # since linear does not support noncontiguous buffer yet
    @dtypes(torch.float, torch.double)
    def test_linear_noncontiguous(self, device, dtype):
        nt_contiguous, nt_noncontiguous = random_nt_noncontiguous_pair(
            (2, 3, 6, 7), device, dtype
        )
        weight = torch.randn((8, 5), device=device, dtype=dtype)
        self.assertRaisesRegex(
            RuntimeError,
            r"for now linear only supports contiguous nested tensor",
            lambda: torch.nn.functional.linear(nt_noncontiguous, weight),
        )

    @dtypes(torch.float, torch.float16, torch.double)
    def test_to_padded_tensor_zero_numel_errors(self, device, dtype):
        ts = [torch.ones(1, 0), torch.ones(0, 0)]
        nt = torch.nested.nested_tensor(
            ts, device=device, dtype=dtype, layout=torch.strided
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"at least one constituent tensor should have non-zero numel",
            lambda: torch.nested.to_padded_tensor(nt, 0.0),
        )

    @dtypes(torch.float, torch.float16, torch.double)
    def test_transpose(self, device, dtype):
        nt = random_nt(device, dtype, 4, (4, 4))
        # error case: transpose nested dimension
        self.assertRaisesRegex(
            RuntimeError,
            "Nested tensor dimension 0 cannot be transposed",
            lambda: nt.transpose(0, 1),
        )
        self.assertRaisesRegex(
            RuntimeError,
            "Nested tensor dimension 0 cannot be transposed",
            lambda: nt.transpose(1, -3),
        )
        # error case: dimension out of range
        self.assertRaises(IndexError, lambda: nt.transpose(1, 3))
        self.assertRaises(IndexError, lambda: nt.transpose(-4, -1))
        # normal case
        ntT = nt.transpose(-1, -2)
        ptT_from_ntT = noncontiguous_to_padded_tensor(ntT)
        pt = torch.nested.to_padded_tensor(nt, 0.0)
        ptT = pt.transpose(-1, -2)
        self.assertEqual(ptT, ptT_from_ntT)

    @dtypes(torch.float, torch.float16, torch.double)
    def test_squeeze_unsqueeze(self, device, dtype):
        a = torch.arange(6).reshape(2, 3)
        b = torch.arange(15).reshape(5, 3)
        nt = torch.nested.nested_tensor([a, b], device=device, dtype=dtype)
        # error case: squeeze no dimension
        self.assertRaisesRegex(
            RuntimeError,
            "For nested tensors, squeeze without the dim argument",
            lambda: nt.squeeze(),
        )
        # error case: squeeze nested dimension
        self.assertRaisesRegex(
            RuntimeError,
            "For nested tensors, squeezing dimension 0",
            lambda: nt.squeeze(0),
        )
        # error case: dimension out of range
        self.assertRaises(IndexError, lambda: nt.squeeze(3))
        # error case: squeeze nested tensor of singleton tensors
        c = torch.ones(1)
        nt_singleton = torch.nested.nested_tensor([c, c], device=device, dtype=dtype)
        self.assertRaisesRegex(
            RuntimeError,
            "For nested tensors, squeezing a nested tensor of singleton",
            lambda: nt_singleton.squeeze(1),
        )

        # squeezing a dim which does not have size 1 should be a no-op
        nt2 = nt.squeeze(-1)
        self.assertEqual(nt, nt2)

        # test cases that should work
        nt_sizes = nt._nested_tensor_size()
        nt_strides = nt._nested_tensor_strides()
        for i in range(-2, 4):
            if i == 0:
                # cannot unsqueeze batch dim
                continue
            nt_unsqueezed = nt.unsqueeze(i)
            # negative dim will correspond to unsqueeze() applied at dim = dim + nt.dim() + 1
            wrapped_i = i + nt.dim() + 1 if i < 0 else i
            # col_index into nt size tensor is requires subtraction of 1 to ignore batch dim
            size_idx = wrapped_i - 1
            self.assertEqual(
                nt_unsqueezed._nested_tensor_size()[:, size_idx],
                torch.ones(2, dtype=torch.long),
            )
            unsqueezed_stride = nt_unsqueezed._nested_tensor_strides()[:, size_idx]
            if i == nt.ndim or i == -1:
                self.assertEqual(unsqueezed_stride, torch.ones(2, dtype=torch.long))
            else:
                stride_col_after = nt_strides[:, size_idx]
                size_col_after = nt_sizes[:, size_idx]
                self.assertEqual(unsqueezed_stride, stride_col_after * size_col_after)
            nt_squeezed = nt_unsqueezed.squeeze(i)
            self.assertEqual(nt_squeezed, nt)
            self.assertEqual(nt_squeezed._nested_tensor_size(), nt_sizes)
            self.assertEqual(nt_squeezed._nested_tensor_strides(), nt_strides)

    @dtypes(torch.float, torch.float16, torch.double)
    def test_transpose_inference_mode_interaction(self, device, dtype):
        nt = random_nt(device, dtype, 4, (4, 4))
        # Construct in default mode and transpose while in inference mode
        with torch.inference_mode():
            ntT = nt.transpose(-1, -2)
            ptT_from_ntT = noncontiguous_to_padded_tensor(ntT)
            pt = torch.nested.to_padded_tensor(nt, 0.0)
            ptT = pt.transpose(-1, -2)
            self.assertEqual(ptT, ptT_from_ntT)

        # Construct and transpose while in inference mode
        with torch.inference_mode():
            nt = random_nt(device, dtype, 4, (4, 4))
            ntT = nt.transpose(-1, -2)
            ptT_from_ntT = noncontiguous_to_padded_tensor(ntT)
            pt = torch.nested.to_padded_tensor(nt, 0.0)
            ptT = pt.transpose(-1, -2)
            self.assertEqual(ptT, ptT_from_ntT)

    @dtypes(torch.float, torch.float16, torch.double)
    def test_view(self, device, dtype):
        nt = random_nt(device, dtype, 4, (4, 4))
        # error case: empty shape
        self.assertRaisesRegex(
            RuntimeError,
            r"shape '\[\]' is invalid for a nested tensor",
            lambda: nt.view(()),
        )
        # error case: empty nested tensor
        nt_empty = torch.nested.nested_tensor([])
        self.assertRaisesRegex(
            RuntimeError,
            "empty nested tensor cannot be reshaped",
            lambda: nt_empty.view(-1),
        )
        # error case: -1 for batch size
        self.assertRaisesRegex(
            RuntimeError,
            r"view: For now nested view cannot change or infer the implicit batch dimension",
            lambda: nt.view(-1, 2, 3),
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"shape '\[.*\]' is invalid for input of size [0-9]+",
            lambda: nt.view(4, 2, 3),
        )
        # normal case
        x0 = torch.randn((2, 20), device=device, dtype=dtype)
        x1 = torch.randn((3, 20), device=device, dtype=dtype)
        nt = torch.nested.nested_tensor([x0, x1])
        pt = torch.nested.to_padded_tensor(nt, 0.0)
        # error case, trying to reshape batch dim to a legit shape
        self.assertRaisesRegex(
            RuntimeError,
            r"For now nested view cannot change or infer the implicit batch dimension",
            lambda: nt.transpose(-1, -2).view(40, -1),
        )
        # inherit only the ragged dimension
        # (2, 20) -> (2, 5, 4)
        # (3, 20) -> (3, 5, 4)
        nt1 = nt.view(2, -1, 5, 4)
        # (2, 3, 20) -> (2, 3, 5, 4) -> (2, 4, 5, 4)
        pt1 = pt.view(2, -1, 5, 4)
        self.assertEqual(noncontiguous_to_padded_tensor(nt1), pt1)

        # more than one -1 (even for "old" dims), should fail
        # this attempts to do # (2, (2, 3), 5, 4) -> (2, (2, 3), 5, 2, 2)
        # but we ban "inherit old behavior" for >1 dimension
        self.assertRaisesRegex(
            RuntimeError,
            r"only one dimension can be inferred",
            lambda: nt1.view(2, -1, -1, 2, 2),
        )

    @dtypes(torch.float, torch.float16, torch.double)
    def test_view_inference_mode_interaction(self, device, dtype):
        # Construct in default mode and view while in inference mode
        nt = torch.nested.nested_tensor(
            [torch.randn((2, 20)), torch.randn((3, 20))], device=device, dtype=dtype
        )
        with torch.inference_mode():
            ntT = nt.view(2, -1, 4, 5)
            ptT_from_ntT = noncontiguous_to_padded_tensor(ntT)
            pt = torch.nested.to_padded_tensor(nt, 0.0)
            ptT = pt.view(2, -1, 4, 5)
            self.assertEqual(ptT, ptT_from_ntT)
        # Construct and view while in inference mode
        with torch.inference_mode():
            nt = torch.nested.nested_tensor(
                [torch.randn((2, 20)), torch.randn((3, 20))], device=device, dtype=dtype
            )
            ntT = nt.view(2, -1, 4, 5)
            ptT_from_ntT = noncontiguous_to_padded_tensor(ntT)
            pt = torch.nested.to_padded_tensor(nt, 0.0)
            ptT = pt.view(2, -1, 4, 5)
            self.assertEqual(ptT, ptT_from_ntT)

    @dtypes(torch.float, torch.float16, torch.double)
    def test_reshape(self, device, dtype):
        nt = random_nt(device, dtype, 4, (4, 4))
        # error case: empty shape
        self.assertRaisesRegex(
            RuntimeError,
            r"shape '\[\]' is invalid for a nested tensor",
            lambda: nt.reshape(()),
        )
        # error case: empty nested tensor
        nt_empty = torch.nested.nested_tensor([])
        self.assertRaisesRegex(
            RuntimeError,
            "empty nested tensor cannot be reshaped",
            lambda: nt_empty.reshape(-1),
        )
        # error case: -1 for batch size
        self.assertRaisesRegex(
            RuntimeError,
            r"reshape: For now nested reshape cannot change or infer the implicit batch dimension",
            lambda: nt.reshape(-1, 2, 3),
        )
        self.assertRaisesRegex(
            RuntimeError,
            r"shape '\[.*\]' is invalid for input of size [0-9]+",
            lambda: nt.reshape(4, 2, 3),
        )
        # normal case
        x0 = torch.randn((2, 20), device=device, dtype=dtype)
        x1 = torch.randn((3, 20), device=device, dtype=dtype)
        nt = torch.nested.nested_tensor([x0, x1])  # (2, (2, 3), 20)
        pt = torch.nested.to_padded_tensor(nt, 0.0)
        # error case, trying to reshape batch dim to a legit shape
        self.assertRaisesRegex(
            RuntimeError,
            r"reshape: For now nested reshape cannot change or infer the implicit batch dimension",
            lambda: nt.transpose(-1, -2).reshape(40, -1),
        )
        # inherit only the ragged dimension
        # (2, 20) -> (2, 5, 4)
        # (3, 20) -> (3, 5, 4)
        nt1 = nt.reshape(2, -1, 5, 4)
        # (2, 3, 20) -> (2, 3, 5, 4) -> (2, 4, 5, 4)
        pt1 = pt.reshape(2, -1, 5, 4)
        self.assertEqual(noncontiguous_to_padded_tensor(nt1), pt1)

        # more than one -1 (even for "old" dims), should fail
        # this attempts to do # (2, (2, 3), 5, 4) -> (2, (2, 3), 5, 2, 2)
        # but we ban "inherit old behavior" for >1 dimension
        self.assertRaisesRegex(
            RuntimeError,
            r"only one dimension can be inferred",
            lambda: nt1.reshape(2, -1, -1, 2, 2),
        )

    def test_nested_masked_select(self, device):
        t = torch.randn([3, 3], device=device)
        mask = torch.tensor([False], device=device)

        njt = torch.nested.masked_select(t, mask)
        self.assertEqual(njt.values(), torch.tensor([], device=device))
        self.assertEqual(njt.offsets(), torch.tensor([0, 0, 0, 0], device=device))

        mask = torch.tensor([[False], [False], [True]], device=device)
        njt = torch.nested.masked_select(t, mask)
        self.assertEqual(njt.values(), t[-1], atol=0.1, rtol=0.1)
        self.assertEqual(njt.offsets(), torch.tensor([0, 0, 0, 3], device=device))

        mask = torch.tensor(
            [[False, False, True], [True, False, True], [False, False, True]],
            device=device,
        )
        njt = torch.nested.masked_select(t, mask)
        self.assertEqual(njt.values(), t.masked_select(mask))
        self.assertEqual(njt.offsets(), torch.tensor([0, 1, 3, 4], device=device))

        t = torch.randn([2, 3, 3, 1], device=device)
        mask = torch.tensor(
            [
                [
                    [[True], [False], [True]],
                    [[True], [False], [True]],
                    [[True], [False], [True]],
                ],
                [
                    [[False], [True], [True]],
                    [[False], [True], [True]],
                    [[True], [True], [True]],
                ],
            ],
            device=device,
        )
        njt = torch.nested.masked_select(t, mask)
        self.assertEqual(njt.values(), t.masked_select(mask))
        self.assertEqual(
            njt.offsets(),
            torch.tensor(
                [0, 1, 1, 2, 3, 3, 4, 5, 5, 6, 6, 7, 8, 8, 9, 10, 11, 12, 13],
                device=device,
            ),
        )

    @dtypes(torch.float, torch.float16, torch.double)
    def test_narrow(self, device, dtype):
        nt = random_nt_from_dims([5, None, None, None], device=device, dtype=dtype)

        # narrow on dim=0 from start to end
        bounds = [(0, 5), (0, 3), (1, 2), (1, 5), (2, 4)]
        for start, end in bounds:
            length = end - start
            narrowed = nt.narrow(dim=0, start=start, length=length)
            # ensure output is a view
            self.assertTrue(narrowed._base is nt)
            for nc, c in zip(narrowed.unbind(), nt.unbind()[start:end]):
                self.assertEqual(nc, c)

        # dim != 0 is not supported
        for dim in range(1, nt.dim()):
            with self.assertRaisesRegex(
                RuntimeError, "only dim=0 supported for nested tensors"
            ):
                nt.narrow(dim=dim, start=0, length=1)

        # error case: non-contiguous NT
        _, nt_noncont = random_nt_noncontiguous_pair((2, 3, 4))
        with self.assertRaisesRegex(
            RuntimeError, "only contiguous nested tensors supported"
        ):
            nt_noncont.narrow(dim=0, start=0, length=1)

    @parametrize("input_dim", [3, 4])
    def test_scaled_dot_product_attention(self, device, input_dim):
        def rand_tensor(*shape):
            return torch.randn(shape, device=device)

        E = 8
        if input_dim == 3:
            # Shape: (N, L, E); ragged L
            query = torch.nested.nested_tensor(
                [rand_tensor(2, E), rand_tensor(3, E), rand_tensor(4, E)]
            )

            # Shape: (N, S, E); ragged S
            key = torch.nested.nested_tensor(
                [rand_tensor(3, E), rand_tensor(4, E), rand_tensor(5, E)]
            )
            value = torch.nested.nested_tensor(
                [rand_tensor(3, E), rand_tensor(4, E), rand_tensor(5, E)]
            )
        elif input_dim == 4:
            # In the 4D case the L and S is ragged
            # Shape: (N, N', L, E); ragged N' and L
            query = torch.nested.nested_tensor(
                [rand_tensor(2, 2, E), rand_tensor(3, 3, E), rand_tensor(4, 4, E)]
            )
            # Shape: (N, N', S, E); ragged N' and S
            key = torch.nested.nested_tensor(
                [rand_tensor(2, 3, E), rand_tensor(3, 4, E), rand_tensor(4, 5, E)]
            )
            value = torch.nested.nested_tensor(
                [rand_tensor(2, 3, E), rand_tensor(3, 4, E), rand_tensor(4, 5, E)]
            )
        else:
            self.fail(f"Invalid input_dim {input_dim} encountered in SDP test")

        def rand_mask(size):
            return torch.randint(0, 2, size=size, dtype=torch.bool, device=device)

        # Shape: (N, L, S); ragged L and S matching above
        attn_mask = torch.nested.nested_tensor(
            [rand_mask((2, 3)), rand_mask((3, 4)), rand_mask((4, 5))]
        )

        dropout_p = 0.0  # no dropout for reproducibility

        # Success case: no attn_mask set and is_causal=False.
        actual = torch.nn.functional.scaled_dot_product_attention(
            query, key, value, attn_mask=None, is_causal=False, dropout_p=dropout_p
        )

        expected_outputs = []
        for q, k, v in zip(query.unbind(), key.unbind(), value.unbind()):
            output = torch.nn.functional.scaled_dot_product_attention(
                q.unsqueeze(0),
                k.unsqueeze(0),
                v.unsqueeze(0),
                attn_mask=None,
                dropout_p=dropout_p,
            )
            expected_outputs.append(output.squeeze(0))
        expected_output_nested = torch.nested.nested_tensor(expected_outputs)
        self.assertEqual(actual, expected_output_nested)

        # Error case: explicit attn_mask set.
        with self.assertRaisesRegex(
            RuntimeError, "not supported when an explicit attn_mask is set"
        ):
            torch.nn.functional.scaled_dot_product_attention(
                query, key, value, attn_mask=attn_mask, dropout_p=dropout_p
            )

        # Error case: is_causal=True.
        with self.assertRaisesRegex(RuntimeError, "not supported when is_causal=True"):
            torch.nn.functional.scaled_dot_product_attention(
                query, key, value, dropout_p=dropout_p, is_causal=True
            )

    @dtypes(torch.float, torch.float16, torch.double)
    def test_empty_like(self, device, dtype):
        ntensors = 4
        nt = random_nt(device, dtype, ntensors, (4, 4))

        # Create empty on same device as original nested tensor
        nt_empty = torch.empty_like(nt)
        assert nt.is_same_size(nt_empty)
        self.assertEqual(nt.dtype, nt_empty.dtype)
        self.assertEqual(nt.device, nt_empty.device)
        self.assertEqual(nt.layout, nt_empty.layout)

        if torch.cuda.is_available():
            if device == "cpu":
                nt_cuda = torch.empty_like(nt, device="cuda")
                self.assertEqual(torch.device("cuda").type, nt_cuda.device.type)
            else:
                nt_cpu = torch.empty_like(nt, device="cpu")
                self.assertEqual(torch.device("cpu").type, nt_cpu.device.type)

        # Check changing dtype of empty_like nested tensor output
        dtype_set = {torch.float, torch.float16, torch.double}
        for other_dtype in dtype_set - {dtype}:
            nt_empty_other_dtype = torch.empty_like(nt, dtype=other_dtype)
            self.assertEqual(nt.dtype, dtype)
            self.assertEqual(nt_empty_other_dtype.dtype, other_dtype)
            self.assertEqual(nt.device, nt_empty.device)
            self.assertEqual(nt.layout, nt_empty.layout)

        # Create tensor for autograd
        nt_empty_req_grad = torch.empty_like(nt, requires_grad=True)
        self.assertEqual(nt_empty_req_grad.requires_grad, True)

        # Test noncontiguous tensor does not fail to copy
        nt_cont, nt_noncont = random_nt_noncontiguous_pair((2, 3, 6, 7))
        nt_empty = torch.empty_like(nt_cont)
        assert nt_cont.is_same_size(nt_empty)
        nt_empty_non_contig = torch.empty_like(nt_noncont)
        assert nt_noncont.is_same_size(nt_empty_non_contig)

        # Test the contiguous memory format option
        nt_empty_contig = torch.empty_like(
            nt_cont, memory_format=torch.contiguous_format
        )
        assert nt_cont.is_same_size(nt_empty_contig)
        assert nt_empty_contig.is_contiguous()

        nt_empty_non_contig = torch.empty_like(
            nt_noncont, memory_format=torch.contiguous_format
        )
        assert nt_noncont.is_same_size(nt_empty_non_contig)
        assert nt_empty_non_contig.is_contiguous()

        # Test other memory formats fail
        self.assertRaises(
            RuntimeError,
            lambda: torch.empty_like(nt_cont, memory_format=torch.channels_last),
        )
        self.assertRaises(
            RuntimeError,
            lambda: torch.empty_like(nt_noncont, memory_format=torch.channels_last),
        )
        self.assertRaises(
            RuntimeError,
            lambda: torch.empty_like(nt_cont, memory_format=torch.channels_last_3d),
        )
        self.assertRaises(
            RuntimeError,
            lambda: torch.empty_like(nt_noncont, memory_format=torch.channels_last_3d),
        )


@markDynamoStrictTest
class TestNestedTensorAutograd(NestedTensorTestCase):
    # Note [Gradcheck args check_batched_grad=False] the common_utils testing version of gradcheck
    # includes the default parameters used for testing ops with gradcheck. However nested tensor
    # does not support the stack op therefore we turn it off for these tests
    def _create_leaf_nested_tensor_from_list(self, tensor_device, requires_grad=False):
        return torch.nested.nested_tensor(
            [torch.randn(1, 2), torch.randn(7, 8)],
            requires_grad=requires_grad,
            device=tensor_device,
        )

    def _create_nested_tensor_from_list(self, tensor_device, requires_grad=False):
        return torch.nested.as_nested_tensor(
            [
                torch.randn(1, 2, requires_grad=requires_grad),
                torch.randn(7, 8, requires_grad=requires_grad),
            ],
            device=tensor_device,
        )

    def _create_nested_tensor_from_mask(self, tensor_device, requires_grad=False):
        data = torch.randn(2, 3, 4, requires_grad=requires_grad, device=tensor_device)
        mask = torch.ones_like(data[:, :, 0]).bool()
        return torch._nested_tensor_from_mask(data, mask)

    def test_as_nested_tensor_propagates_gradients(self, device):
        a = torch.arange(3, dtype=torch.float, device=device)
        b = torch.arange(5, dtype=torch.float, device=device)
        nt = torch.nested.as_nested_tensor([a, b])
        # tensors with requires_grad=False are leaves
        self.assertTrue(nt.is_leaf)
        self.assertTrue(not nt.requires_grad)

        a = torch.arange(3, dtype=torch.float, requires_grad=True, device=device)
        b = torch.arange(5, dtype=torch.float, requires_grad=True, device=device)
        nt2 = torch.nested.as_nested_tensor([a, b])
        fake_grad = torch.nested.nested_tensor(
            [torch.ones_like(a), torch.zeros_like(b)], device=device
        )
        nt2.backward(fake_grad)
        self.assertEqual(a.grad, fake_grad[0])
        self.assertEqual(b.grad, fake_grad[1])

    def test_nested_tensor_generates_leaf(self, device):
        a = torch.arange(3, dtype=torch.float, requires_grad=True, device=device)
        b = torch.arange(5, dtype=torch.float, requires_grad=True, device=device)

        nt = torch.nested.nested_tensor([a, b], requires_grad=False)
        self.assertTrue(nt.is_leaf)
        self.assertTrue(not nt.requires_grad)

        nt2 = torch.nested.nested_tensor([a, b], requires_grad=True)
        self.assertTrue(nt2.is_leaf)
        self.assertTrue(nt2.requires_grad)

        fake_grad = torch.nested.nested_tensor(
            [torch.ones_like(a), torch.zeros_like(b)], device=device
        )
        nt2.backward(fake_grad)
        self.assertEqual(nt2.grad, fake_grad)
        self.assertEqual(a.grad, None)
        self.assertEqual(b.grad, None)

    def test_set_requires_grad_from_list(self, device):
        nt = self._create_nested_tensor_from_list(device)
        nt.requires_grad_()
        assert nt.requires_grad

    def test_set_requires_grad_from_mask(self, device):
        nt = self._create_nested_tensor_from_mask(device)
        nt.requires_grad_()
        assert nt.requires_grad

    def test_backward_for_add_op(self, device):
        nt_1 = self._create_nested_tensor_from_mask(device)
        nt_2 = self._create_nested_tensor_from_mask(device)

        nt_1.requires_grad_()
        c = nt_1 + nt_2

        assert nt_1.requires_grad
        assert c.requires_grad
        grad_output = self._create_nested_tensor_from_mask(device)
        c.backward(grad_output)

        #  Grad check doesn't work with nested yet.
        # d/dnt_1 (nt + nt_1) = 1*grad_output
        self.assertEqual(nt_1.grad, grad_output)

    def test_backward_for_sub_op(self, device):
        nt_1 = self._create_nested_tensor_from_mask(device)
        nt_2 = self._create_nested_tensor_from_mask(device)

        nt_1.requires_grad_()
        nt_2.requires_grad_()
        c = nt_1 - nt_2

        assert nt_1.requires_grad
        assert nt_2.requires_grad
        assert c.requires_grad
        grad_output = self._create_nested_tensor_from_mask(device)
        c.backward(grad_output)

        self.assertEqual(nt_1.grad, grad_output)
        self.assertEqual(nt_2.grad, -1 * grad_output)

    def test_backward_sub_strided(self, device):
        a = torch.nested.nested_tensor(
            [torch.randn(9, 2, 4), torch.randn(12, 2, 4)],
            requires_grad=True,
            device=device,
        )
        b = torch.nested.nested_tensor(
            [torch.randn(9, 4, 2), torch.randn(12, 4, 2)],
            requires_grad=True,
            device=device,
        )
        c = a - b.transpose(-1, -2)
        grad_output = c.clone()
        c.backward(grad_output)
        self.assertEqual(a.grad, grad_output)
        self.assertEqual(b.grad, -1 * grad_output.transpose(-1, -2))

    def test_backward_add_strided(self, device):
        a = torch.nested.nested_tensor(
            [torch.randn(9, 2, 4), torch.randn(12, 2, 4)],
            requires_grad=True,
            device=device,
        )
        b = torch.nested.nested_tensor(
            [torch.randn(9, 4, 2), torch.randn(12, 4, 2)],
            requires_grad=True,
            device=device,
        )
        c = a + b.transpose(-1, -2)
        grad_output = c.clone()
        c.backward(grad_output)
        self.assertEqual(a.grad, grad_output)
        self.assertEqual(b.grad, grad_output.transpose(-1, -2))

    # Test Factory Functions
    def test_nested_tensor_to_padded_tensor(self, device):
        for padding_val in [0, 1]:
            nt = self._create_leaf_nested_tensor_from_list(
                tensor_device=device, requires_grad=True
            )

            out = torch.nested.to_padded_tensor(nt, padding_val)
            grad_output = torch.ones(out.shape, device=device)
            out.backward(grad_output)

            self.assertEqual(
                nt.grad,
                torch.nested.nested_tensor(
                    [torch.ones(1, 2), torch.ones(7, 8)], device=device
                ),
            )

    def test_nested_tensor_from_mask_and_to_padded(self, device):
        N, L, D = 2, 4, 4
        mask = torch.ones(N, L, device=device)
        for i in range(1, N):
            end = torch.randint(1, L - 1, (1,), device=device)
            mask[i, end:] = 0

        mask[0, :] = 1
        mask = mask.bool()

        data = torch.randn(
            N, L, D, requires_grad=True, dtype=torch.float64, device=device
        )

        def grad_test_func(inpt):
            nt = torch._nested_tensor_from_mask(inpt, mask)
            # This implicitly tests to_padded_tensor grads
            return torch.nested.to_padded_tensor(nt, 0)

        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_nested_tensor_from_padded(self, device):
        nested_size = torch.tensor([[1, 2], [2, 2]])
        padded_tensor = torch.randn(2, 2, 2, dtype=torch.float64, device=device)
        padded_tensor[0, 1, :] = 0
        padded_tensor.requires_grad_()

        def grad_test_func(tensor, nested_size):
            nt = torch._nested_from_padded(
                tensor, nested_size, fuse_transform_0213=False
            )
            # This implicitly tests to_padded_tensor grads
            return torch.nested.to_padded_tensor(nt, 0)

        data = (padded_tensor, nested_size)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_nested_tensor_from_padded_fused(self, device):
        nested_size = torch.tensor([[1, 8], [2, 8]])
        padded_tensor = torch.randn(2, 2, 2, 4, dtype=torch.float64, device=device)
        padded_tensor[0, 1, :] = 0
        padded_tensor.requires_grad_()

        def grad_test_func(tensor, nested_size):
            nt = torch._nested_from_padded(
                tensor, nested_size, fuse_transform_0213=True
            )
            # This implicitly tests to_padded_tensor grads
            return torch.nested.to_padded_tensor(nt, 0)

        data = (padded_tensor, nested_size)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_nested_tensor_from_list(self, device):
        a = torch.randn(1, 2, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(10, 2, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            c = torch.nested.as_nested_tensor([a, b, c])
            # This implictily tests to_padded_tensor grads
            return torch.nested.to_padded_tensor(c, 0)

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    @parametrize("layout", [torch.strided, torch.jagged], name_fn=layout_name)
    def test_dropout_backward(self, layout):
        if layout == torch.jagged:
            nt = torch.nested.nested_tensor(
                [torch.randn((2, 5)), torch.randn((3, 5))],
                requires_grad=True,
                layout=layout,
            )
        else:
            nt = torch.nested.nested_tensor(
                [torch.randn((2, 5)), torch.randn((3, 4))],
                requires_grad=True,
                layout=layout,
            )
        p = 0.2
        y = torch.nn.functional.dropout(nt, p)
        y.backward(nt.clone().detach())
        self.assertEqual(nt.grad, y)

    def test_nested_tensor_bmm_gradcheck(self, device):
        a = torch.randn(2, 6, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(3, 6, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(6, 4, requires_grad=True, dtype=torch.float64, device=device)
        d = torch.randn(6, 5, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c, d):
            nt0 = torch.nested.as_nested_tensor([a, b])
            nt1 = torch.nested.as_nested_tensor([c, d])
            result = nt0.bmm(nt1)
            return torch.nested.to_padded_tensor(result, 0.0)

        data = (a, b, c, d)
        assert torch.autograd.gradcheck(grad_test_func, inputs=data)

    def test_nested_tensor_bmm_backward(self, device):
        nt0 = torch.nested.nested_tensor(
            [torch.randn((2, 6)), torch.randn((3, 6))],
            requires_grad=True,
            device=device,
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((6, 4)), torch.randn((6, 5))],
            requires_grad=True,
            device=device,
        )
        with torch.no_grad():
            pt0 = torch.nested.to_padded_tensor(nt0, 0.0).requires_grad_(True)
            pt1 = torch.nested.to_padded_tensor(nt1, 0.0).requires_grad_(True)

        ynt = nt0.bmm(nt1)
        ypt = pt0.bmm(pt1)
        ynt.backward(ynt.clone())
        ypt.backward(ypt.clone())

        self.assertEqual(torch.nested.to_padded_tensor(nt0.grad, 0.0), pt0.grad)
        self.assertEqual(torch.nested.to_padded_tensor(nt1.grad, 0.0), pt1.grad)

    def test_nested_tensor_matmul_gradcheck(self, device):
        a = torch.randn(2, 6, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(3, 6, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(6, 4, requires_grad=True, dtype=torch.float64, device=device)
        d = torch.randn(6, 5, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c, d):
            nt0 = torch.nested.as_nested_tensor([a, b])
            nt1 = torch.nested.as_nested_tensor([c, d])
            result = torch.matmul(nt0, nt1)
            return torch.nested.to_padded_tensor(result, 0.0)

        data = (a, b, c, d)
        assert torch.autograd.gradcheck(grad_test_func, inputs=data)

    def test_nested_tensor_matmul_backward(self, device):
        nt0 = torch.nested.nested_tensor(
            [torch.randn((7, 2, 6)), torch.randn((7, 3, 6))],
            requires_grad=True,
            device=device,
        )
        nt1 = torch.nested.nested_tensor(
            [torch.randn((7, 6, 4)), torch.randn((7, 6, 5))],
            requires_grad=True,
            device=device,
        )
        with torch.no_grad():
            pt0 = torch.nested.to_padded_tensor(nt0, 0.0).requires_grad_(True)
            pt1 = torch.nested.to_padded_tensor(nt1, 0.0).requires_grad_(True)

        ynt = torch.matmul(nt0, nt1)
        ypt = torch.matmul(pt0, pt1)
        ynt.backward(ynt.clone())
        ypt.backward(ypt.clone())

        self.assertEqual(torch.nested.to_padded_tensor(nt0.grad, 0.0), pt0.grad)
        self.assertEqual(torch.nested.to_padded_tensor(nt1.grad, 0.0), pt1.grad)

    def test_nested_tensor_transpose_gradcheck(self, device):
        a = torch.randn(2, 5, requires_grad=True, device=device)
        b = torch.randn(3, 4, requires_grad=True, device=device)

        def grad_test_func(a, b):
            nt = torch.nested.as_nested_tensor([a, b])
            result = nt.transpose(-2, -1).transpose(-2, -1)
            return torch.nested.to_padded_tensor(result, 0.0)

        data = (a, b)
        assert torch.autograd.gradcheck(grad_test_func, inputs=data, eps=1e-3)

    def test_nested_tensor_transpose_backward(self, device):
        nt = torch.nested.nested_tensor(
            [torch.randn((2, 5)), torch.randn((3, 4))],
            requires_grad=True,
            device=device,
        )
        with torch.no_grad():
            pt = torch.nested.to_padded_tensor(nt, 0.0).requires_grad_(True)

        ynt = nt.transpose(-2, -1)
        ypt = pt.transpose(-2, -1)
        ynt.backward(ynt.clone())
        ypt.backward(ypt.clone())

        self.assertEqual(torch.nested.to_padded_tensor(nt.grad, 0.0), pt.grad)

    def test_nested_tensor_reshape_gradcheck(self, device):
        a = torch.randn(2, 6, requires_grad=True, device=device)
        b = torch.randn(3, 6, requires_grad=True, device=device)

        def grad_test_func(a, b):
            nt = torch.nested.as_nested_tensor([a, b])
            result = nt.reshape(2, -1, 2, 3)
            return torch.nested.to_padded_tensor(result, 0.0)

        data = (a, b)
        assert torch.autograd.gradcheck(grad_test_func, inputs=data, eps=1e-3)

    def test_nested_tensor_reshape_backward(self):
        nt = torch.nested.nested_tensor(
            [torch.randn((2, 6)), torch.randn((3, 6))], requires_grad=True
        )
        with torch.no_grad():
            pt = torch.nested.to_padded_tensor(nt, 0.0).requires_grad_(True)

        ynt = nt.reshape(2, -1, 2, 3)
        ypt = pt.reshape(2, -1, 2, 3)
        ynt.backward(ynt.clone())
        ypt.backward(ypt.clone())

        self.assertEqual(torch.nested.to_padded_tensor(nt.grad, 0.0), pt.grad)

    def test_nested_tensor_squeeze_backward(self, device):
        nt = torch.nested.nested_tensor(
            [torch.randn((2, 6, 1)), torch.randn((3, 6, 1))],
            requires_grad=True,
            device=device,
        )
        with torch.no_grad():
            pt = torch.nested.to_padded_tensor(nt, 0.0).requires_grad_(True)

        ynt = nt.squeeze(-1)
        ypt = pt.squeeze(-1)
        ynt.backward(ynt.clone())
        ypt.backward(ypt.clone())

        self.assertEqual(torch.nested.to_padded_tensor(nt.grad, 0.0), pt.grad)

    def test_nested_tensor_squeeze_gradcheck(self, device):
        a = torch.randn(
            (2, 6, 1), dtype=torch.float64, requires_grad=True, device=device
        )
        b = torch.randn(
            (3, 6, 1), dtype=torch.float64, requires_grad=True, device=device
        )

        def grad_test_func(a, b):
            nt = torch.nested.as_nested_tensor([a, b])
            result = nt.squeeze(-1)
            return torch.nested.to_padded_tensor(result, 0.0)

        assert torch.autograd.gradcheck(grad_test_func, inputs=(a, b), eps=1e-3)

    def test_nested_tensor_unsqueeze_backward(self, device):
        nt = torch.nested.nested_tensor(
            [torch.randn((2, 6)), torch.randn((3, 6))],
            requires_grad=True,
            device=device,
        )
        with torch.no_grad():
            pt = torch.nested.to_padded_tensor(nt, 0.0).requires_grad_(True)

        ynt = nt.unsqueeze(2)
        ypt = pt.unsqueeze(2)
        ynt.backward(ynt.clone())
        ypt.backward(ypt.clone())

        self.assertEqual(torch.nested.to_padded_tensor(nt.grad, 0.0), pt.grad)

    def test_nested_tensor_unsqueeze_gradcheck(self, device):
        a = torch.randn((2, 6), dtype=torch.float64, requires_grad=True, device=device)
        b = torch.randn((3, 6), dtype=torch.float64, requires_grad=True, device=device)

        def grad_test_func(a, b):
            nt = torch.nested.as_nested_tensor([a, b])
            result = nt.unsqueeze(-1)
            return torch.nested.to_padded_tensor(result, 0.0)

        assert torch.autograd.gradcheck(grad_test_func, inputs=(a, b), eps=1e-3)

    def test_nested_tensor_linear(self, device):
        a = torch.randn(1, 2, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(3, 2, requires_grad=True, dtype=torch.float64, device=device)

        weight = torch.randn(
            2, 2, requires_grad=True, dtype=torch.float64, device=device
        )
        bias = torch.randn(2, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c, weight, bias=None):
            nt = torch.nested.as_nested_tensor([a, b, c])
            # This implicitly tests to_padded_tensor grads
            d = torch.functional.F.linear(nt, weight, bias)
            return torch.nested.to_padded_tensor(d, 0)

        data = (a, b, c, weight, bias)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

        # Test linear with no bias added
        data = (a, b, c, weight)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_nested_tensor_linear_plus_transpose(self, device):
        a = torch.randn(1, 2, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(3, 2, requires_grad=True, dtype=torch.float64, device=device)

        weight = torch.randn(
            2, 2, requires_grad=True, dtype=torch.float64, device=device
        )
        bias = torch.randn(2, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c, weight, bias=None):
            nt = torch.nested.as_nested_tensor([a, b, c])
            # This implicitly tests to_padded_tensor grads
            d = torch.functional.F.linear(nt, weight, bias)
            d = d.transpose(-1, -2).contiguous()
            return torch.nested.to_padded_tensor(d, 0)

        data = (a, b, c, weight, bias)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

        # Test linear with no bias added
        data = (a, b, c, weight)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_nested_tensor_softmax(self, device):
        a = torch.randn(1, 2, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(3, 2, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c, dim):
            nt = torch.nested.as_nested_tensor([a, b, c])
            # This implicitly tests to_padded_tensor grads
            d = torch.functional.F.softmax(nt, dim=dim)
            return torch.nested.to_padded_tensor(d, 0)

        # softmax over last dim
        data = (a, b, c, -1)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_nested_tensor_linear_backward(self, device):
        a = torch.randn(1, 2, requires_grad=False, device=device)
        b = torch.randn(2, 2, requires_grad=False, device=device)
        c = torch.randn(3, 2, requires_grad=False, device=device)

        weight = torch.randn(2, 2, requires_grad=True, device=device)
        bias = torch.randn(2, requires_grad=True, device=device)
        nt = torch.nested.as_nested_tensor([a, b, c], device=device)

        out = torch.functional.F.linear(nt, weight, bias)

        out.backward(out.clone())

        assert weight.grad is not None
        assert bias.grad is not None

        assert a.grad is None
        assert b.grad is None
        assert c.grad is None

    def test_values_grad_with_broadcast(self, device):
        a = torch.randn(1, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(3, 2, 4, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            buffer = nt.values()
            return buffer.sum()

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_to_buffer_series_ops_grad_with_broadcast(self, device):
        a = torch.randn(1, 1, 2, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(1, 1, 2, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(1, 1, 2, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            buffer = nt.values()
            buffer = buffer * 2
            return buffer.exp()

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_unbind_flow_through(self, device):
        a = torch.randn(1, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(3, 2, 4, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            ntT = nt.transpose(-1, -2)
            unbound = ntT.unbind()
            d = unbound[0]
            d = torch.pow(d, 2)
            return d

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_split_with_sizes_flow_through(self, device):
        a = torch.randn(2, 5, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(3, 5, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(4, 5, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            splits = nt.split_with_sizes([2, 3], dim=-1)
            unbound = splits[1].unbind()
            d = unbound[0]
            d = torch.pow(d, 2)
            return d

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_indexing_backward(self, device):
        x0 = torch.randn((2, 5))
        x1 = torch.randn((3, 4))
        nt = torch.nested.nested_tensor([x0, x1], device=device, requires_grad=True)
        self.assertEqual(nt[0], x0)
        self.assertEqual(nt[-1], x1)
        grad_x0 = torch.randn((2, 5), device=device)
        nt[0].backward(grad_x0)
        expected_grad = torch.nested.nested_tensor(
            [grad_x0, torch.zeros((3, 4), device=device)]
        )
        self.assertEqual(nt.grad, expected_grad)

    def test_masked_fill_backward(self, device):
        a = torch.randn(1, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(3, 2, 4, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            mask = nt.detach().clone().to(bool)
            out = nt.masked_fill(mask, 0)
            out = torch.nested.to_padded_tensor(out, 0)
            return out

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_gelu_backward(self, device):
        a = torch.randn(1, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(3, 2, 4, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            nt_gelu = torch.nn.functional.gelu(nt)
            return torch.nested.to_padded_tensor(nt_gelu, 0)

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_relu_backward(self, device):
        a = torch.randn(1, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(3, 2, 4, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            nt_relu = torch.nn.functional.relu(nt)
            return torch.nested.to_padded_tensor(nt_relu, 0)

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_selu_backward(self, device):
        a = torch.randn(1, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(3, 2, 4, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            nt_relu = torch.nn.functional.silu(nt)
            return torch.nested.to_padded_tensor(nt_relu, 0)

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    def test_abs_backward(self, device):
        a = torch.randn(1, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(2, 2, 4, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(3, 2, 4, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            nt_abs = torch.abs(nt)
            return torch.nested.to_padded_tensor(nt_abs, 0)

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    # Previously would error when input NT doesn't require grad
    # NotImplementedError: Cannot access storage of UndefinedTensorImpl
    def test_layer_norm_backward_edge_case(self, device):
        size = 4
        a = torch.randn(
            1, 2, size, requires_grad=False, dtype=torch.float64, device=device
        )
        nt = torch.nested.nested_tensor([a])
        nt_layer_norm = torch.nn.LayerNorm(
            nt.size(-1), device=device, dtype=torch.float64
        )
        out = nt_layer_norm(nt)
        out.backward(out.clone())

    def test_accumulate_grad_different_strides(self, device):
        a = torch.rand(1, 4, 2, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.rand(1, 8, 2, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b):
            nt_1 = torch.nested.as_nested_tensor([a, b])
            nt_2 = nt_1.clone()
            out = torch.nn.functional.scaled_dot_product_attention(nt_1, nt_2, nt_2)
            return torch.nested.to_padded_tensor(out, 0)

        data = (a, b)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    # https://github.com/pytorch/pytorch/issues/95562
    @skipIfSlowGradcheckEnv
    @parametrize("size", [1024, 1023, 513, 512, 256, 128, 32, 4, 2])
    def test_layer_norm_backward(self, device, size):
        a = torch.randn(
            1, 2, size, requires_grad=True, dtype=torch.float64, device=device
        )
        b = torch.randn(
            2, 2, size, requires_grad=True, dtype=torch.float64, device=device
        )
        c = torch.randn(
            3, 2, size, requires_grad=True, dtype=torch.float64, device=device
        )

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            layer_norm = torch.nn.LayerNorm(
                nt.size(-1), device=device, dtype=torch.float64
            )
            nt_layer_norm = layer_norm(nt)
            return torch.nested.to_padded_tensor(nt_layer_norm, 0)

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)

    # https://github.com/pytorch/pytorch/issues/95562
    @skipIfSlowGradcheckEnv
    # Could either mark slow or reduce size
    @parametrize("size", [128, 32, 4, 2])
    def test_layer_norm_backward_5d(self, device, size):
        a = torch.randn(
            4, size, size, 4, requires_grad=True, dtype=torch.float64, device=device
        )
        b = torch.randn(
            7, size, size, 4, requires_grad=True, dtype=torch.float64, device=device
        )
        c = torch.randn(
            10, size, size, 4, requires_grad=True, dtype=torch.float64, device=device
        )

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c])
            layer_norm = torch.nn.LayerNorm(
                (size, size, nt.size(-1)), device=device, dtype=torch.float64
            )
            nt_layer_norm = layer_norm(nt)
            return torch.nested.to_padded_tensor(nt_layer_norm, 0)

        data = (a, b, c)
        assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)


# Found in torch/testing/_comparison.py
default_atol = {torch.float16: 1e-3, torch.bfloat16: 1e-3, torch.float32: 1e-5}
default_rtol = {torch.float16: 1e-3, torch.bfloat16: 1.6e-2, torch.float32: 1.3e-6}


def get_rtol(true_value: torch.Tensor, computed_value: torch.Tensor) -> float:
    deviation = true_value - computed_value
    deviation = torch.abs(deviation / true_value)
    # Fill in the nans with the default rtol
    torch.nan_to_num_(deviation, nan=default_rtol[computed_value.dtype])
    return deviation.max().item()


def get_atol(true_value: torch.Tensor, computed_value: torch.Tensor) -> float:
    deviation = true_value - computed_value
    atol = torch.abs(deviation).max().item()
    return atol


def get_tolerances(
    true_value: torch.Tensor,
    computed_value: torch.Tensor,
    fudge_factor: Optional[float] = None,
) -> Tuple[float, float]:
    """Returns the absolute and relative tolerances for comparing two tensors."""
    fudge_factor = fudge_factor if fudge_factor is not None else 1.0
    atol = get_atol(true_value, computed_value)
    rtol = get_rtol(true_value, computed_value)

    atol = fudge_factor * max(atol, default_atol[computed_value.dtype])
    rtol = fudge_factor * max(rtol, default_rtol[computed_value.dtype])
    # torch.isclose() has weird behavior around see:
    # https://github.com/pytorch/pytorch/issues/102400
    if rtol > 1e30:
        rtol = default_rtol[computed_value.dtype]
    return atol, rtol


# We can probably parametrizing existing tests instead of having a separate
# test class as we begin to support more ops. Also maybe rewrite with OpInfos.
@markDynamoStrictTest
class TestNestedTensorSubclass(NestedTensorTestCase):
    # TODO: consolidate with the below
    def _get_list_for_jagged_tensor(self, nested_size, device, requires_grad=True):
        Ds = nested_size[1:]
        out = []
        for s in nested_size[0]:
            out.append(
                torch.randn(
                    s,
                    *Ds,
                    requires_grad=requires_grad,
                    device=device,
                    dtype=torch.float64,
                )
            )
        return out

    def _get_example_tensor_lists(
        self,
        include_list_of_lists=True,
        include_requires_grad=True,
        include_inner_dim_size_1=False,
        include_2d_tensor=False,
    ):
        def _make_tensor(
            *shape, include_requires_grad=include_requires_grad, requires_grad=True
        ):
            return torch.randn(
                *shape,
                requires_grad=(requires_grad if include_requires_grad else False),
            )

        # Purposefully introduce mixed requires_grad settings for the components
        # when include_requires_grad=True.
        example_lists = [
            # (B, *, D) with B=4
            [
                _make_tensor(2, 5),
                _make_tensor(3, 5, requires_grad=False),
                _make_tensor(4, 5, requires_grad=False),
                _make_tensor(6, 5),
            ],
            # (B, *, D_0, D_1) with B=5
            [
                _make_tensor(2, 5, 6),
                _make_tensor(3, 5, 6),
                _make_tensor(4, 5, 6, requires_grad=False),
                _make_tensor(5, 5, 6),
                _make_tensor(6, 5, 6),
            ],
            # (B, *, D_0, D_1, D_2) with B=6
            [
                _make_tensor(2, 5, 6, 7),
                _make_tensor(3, 5, 6, 7),
                _make_tensor(4, 5, 6, 7, requires_grad=False),
                _make_tensor(5, 5, 6, 7),
                _make_tensor(6, 5, 6, 7),
                _make_tensor(7, 5, 6, 7),
            ],
        ]

        if include_list_of_lists:
            example_lists.append(
                # (B, *, D) with B=3 in list form
                [
                    _make_tensor(2, 5, requires_grad=False).tolist(),
                    _make_tensor(3, 5).tolist(),
                    _make_tensor(4, 5).tolist(),
                ]
            )

        if include_inner_dim_size_1:
            example_lists.append(
                [
                    _make_tensor(2, 1),
                    _make_tensor(3, 1, requires_grad=False),
                    _make_tensor(4, 1, requires_grad=False),
                    _make_tensor(6, 1),
                ]  # (B, *, 1)
            )
            example_lists.append(
                [
                    _make_tensor(2, 5, 1),
                    _make_tensor(3, 5, 1, requires_grad=False),
                    _make_tensor(4, 5, 1, requires_grad=False),
                    _make_tensor(6, 5, 1),
                ]  # (B, *, 5, 1)
            )

        if include_2d_tensor:
            example_lists.append(
                [
                    _make_tensor(2),
                    _make_tensor(3, requires_grad=False),
                    _make_tensor(4, requires_grad=False),
                    _make_tensor(6),
                ]  # (B, *)
            )

        return example_lists

    def test_tensor_attributes(self, device):
        a = torch.randn(2, 3, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(3, 3, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(4, 3, requires_grad=True, dtype=torch.float64, device=device)
        nt = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)
        _offsets = nt.offsets()

        for op in (
            torch.ops.aten.is_non_overlapping_and_dense.default,
            torch.ops.aten.sym_size.default,
            torch.ops.aten.dim.default,
            torch.ops.aten.numel.default,
            torch.ops.aten.sym_numel.default,
            torch.ops.aten.sym_stride.default,
            torch.ops.aten.sym_storage_offset.default,
        ):
            op(nt)

        with self.assertRaisesRegex(
            RuntimeError, "directly calling torch.ops.aten.size"
        ):
            torch.ops.aten.size.default(nt)

        nested_int = torch.nested._internal.nested_tensor.get_tensor_symint(
            _offsets, coeff=1
        )
        self.assertEqual(nt.size(), (3, nested_int, 3))
        self.assertEqual(nt.shape, (3, nested_int, 3))
        self.assertEqual(nt.dim(), 3)
        self.assertEqual(nt.numel(), 27)

    @parametrize("nt_dim", [3, 4, 5])
    def test_linear(self, device, nt_dim):
        if nt_dim == 3:
            fixed_shape = (3,)
        elif nt_dim == 4:
            fixed_shape = (4, 3)
        elif nt_dim == 5:
            fixed_shape = (5, 4, 3)

        a = torch.randn(
            2, *fixed_shape, requires_grad=True, dtype=torch.float64, device=device
        )
        b = torch.randn(
            3, *fixed_shape, requires_grad=True, dtype=torch.float64, device=device
        )
        c = torch.randn(
            4, *fixed_shape, requires_grad=True, dtype=torch.float64, device=device
        )
        weight = torch.randn(
            4, 3, requires_grad=True, dtype=torch.float64, device=device
        )

        def grad_test_func(a, b, c, weight):
            nt = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)
            out = torch.nn.functional.linear(nt, weight)
            return out.values()

        gradcheck(grad_test_func, inputs=(a, b, c, weight), check_batched_grad=False)

    def test_unary_pointwise(self, device):
        a = torch.randn(2, 3, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(3, 3, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(4, 3, requires_grad=True, dtype=torch.float64, device=device)

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)
            out = torch.nn.functional.silu(nt.sin().cos())
            return out.values()

        gradcheck(grad_test_func, inputs=(a, b, c), check_batched_grad=False)

    def test_unary_pointwise_transposed_inputs(self, device):
        a, b, c = (
            torch.randn(
                i + 2, 5, requires_grad=True, dtype=torch.float64, device=device
            )
            for i in range(3)
        )

        nt = torch.nested.nested_tensor(
            [a.detach(), b.detach(), c.detach()], layout=torch.jagged
        )
        nt_t = nt.transpose(1, 2)
        self.assertFalse(nt_t.is_contiguous())
        out = torch.nn.functional.silu(nt_t.sin().cos())
        self.assertEqual(
            out.is_contiguous(),
            torch.nn.functional.silu(b.transpose(-1, -2).sin().cos()).is_contiguous(),
        )

        self.assertEqual(nt_t.shape, out.shape)

        a, b, c = (
            torch.randn(
                i + 2, 5, requires_grad=True, dtype=torch.float64, device=device
            )
            for i in range(3)
        )

        def grad_test_func(a, b, c):
            nt = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)
            nt_t = nt.transpose(1, 2)
            out = torch.nn.functional.silu(nt_t.sin().cos())
            return out.values()

        gradcheck(grad_test_func, inputs=(a, b, c), check_batched_grad=False)

    def test_binary_pointwise(self, device):
        a = torch.randn(2, 3, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(3, 3, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(4, 3, requires_grad=True, dtype=torch.float64, device=device)

        # Incorrect usage: shape check will fail if the offsets tensor are not
        #                  the same exact tensor object
        nt1 = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)
        nt2 = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)

        self.assertRaisesRegex(
            RuntimeError,
            "cannot call binary pointwise function .* with inputs of shapes",
            lambda: nt1 * nt2,
        )

        # Correct usage: chain the calls using the same offsets tensor object
        def grad_test_func(a, b, c):
            nt1 = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)
            # TODO: Switch to public API that takes in (values, offsets) once it exists
            nt2, offsets = jagged_from_list([a, b, c], nt1.offsets())
            out = nt1 * nt2
            return out.values()

        gradcheck(grad_test_func, inputs=(a, b, c), check_batched_grad=False)

    def test_binary_pointwise_transposed(self, device):
        a, b, c = (
            torch.randn(i + 2, 5, dtype=torch.float64, device=device) for i in range(3)
        )

        nt1, offsets = jagged_from_list([a, b, c], None)
        nt2, offsets = jagged_from_list([a, b, c], offsets)

        nt1_t = nt1.transpose(1, 2)
        nt2_t = nt2.transpose(1, 2)

        # out = nt1_t * nt2_t
        # self.assertFalse(nt1_t.is_contiguous())
        # self.assertEqual(out.is_contiguous(), (b.transpose(-1, -2) * b.transpose(-1, -2)).is_contiguous())
        # self.assertEqual(out.shape, nt1_t.shape)

        self.assertRaisesRegex(
            RuntimeError,
            "cannot call binary pointwise function mul.Tensor with inputs of shapes",
            lambda: nt1 * nt2_t,
        )

        a, b, c = (
            torch.randn(
                i + 2, 5, requires_grad=True, dtype=torch.float64, device=device
            )
            for i in range(3)
        )

        # Correct usage: chain the calls using the same offsets tensor object
        def grad_test_func(a, b, c):
            nt1, offsets = jagged_from_list([a, b, c], None)
            nt2, offsets = jagged_from_list([a, b, c], offsets)
            nt1_t = nt1.transpose(1, 2)
            nt2_t = nt2.transpose(1, 2)
            out = nt1_t * nt2_t
            return out.values()

        gradcheck(grad_test_func, inputs=(a, b, c), check_batched_grad=False)

    def test_split(self, device):
        a = torch.randn(2, 3, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(3, 3, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(4, 3, requires_grad=True, dtype=torch.float64, device=device)

        nt = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)
        out = torch.split(nt, 2, -1)
        self.assertEqual(len(out), 2)
        self.assertEqualIgnoringNestedInts(
            out[0],
            torch.nested.as_nested_tensor(
                [a[:, 0:2], b[:, 0:2], c[:, 0:2]], layout=torch.jagged
            ),
        )
        self.assertEqualIgnoringNestedInts(
            out[1],
            torch.nested.as_nested_tensor(
                [a[:, 2:], b[:, 2:], c[:, 2:]], layout=torch.jagged
            ),
        )

        with self.assertRaisesRegex(
            RuntimeError,
            r"split\(\): not supported for NestedTensor on dim=1",
        ):
            torch.split(nt, 2, 1)

    def test_split_with_sizes(self, device):
        a = torch.randn(2, 3, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(3, 3, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(4, 3, requires_grad=True, dtype=torch.float64, device=device)

        nt = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)
        out = torch.split(nt, [1, 2], -1)
        self.assertEqual(len(out), 2)
        self.assertEqualIgnoringNestedInts(
            out[0],
            torch.nested.as_nested_tensor(
                [a[:, 0:1], b[:, 0:1], c[:, 0:1]], layout=torch.jagged
            ),
        )
        self.assertEqualIgnoringNestedInts(
            out[1],
            torch.nested.as_nested_tensor(
                [a[:, 1:], b[:, 1:], c[:, 1:]], layout=torch.jagged
            ),
        )
        with self.assertRaisesRegex(
            RuntimeError,
            r"split_with_sizes\(\): not supported for NestedTensor on dim=1",
        ):
            torch.split(nt, [1, 2], 1)

    def test_softmax(self, device):
        nt = random_nt_from_dims(
            [3, None, 5],
            device=device,
            dtype=torch.float32,
            layout=torch.jagged,
            requires_grad=True,
        )

        # operate on dim=2
        output = nt.softmax(dim=2)

        @torch._dynamo.disable
        def _compare_to_ref(nt, output, dim):
            for in_component, out_component in zip(nt.unbind(), output.unbind()):
                self.assertEqual(in_component.softmax(dim=dim), out_component)

        # dim=2 -> dim=1 after unbind
        _compare_to_ref(nt, output, dim=1)

        # operate on dim=-1
        output2 = nt.softmax(dim=-1)
        torch._dynamo.disable(self.assertEqual)(output, output2)
        _compare_to_ref(nt, output2, dim=-1)

        def grad_test_func(a, b):
            nt = torch.nested.as_nested_tensor([a, b], layout=torch.jagged)
            out = nt.softmax(dim=-1)
            return out.values()

        a = torch.rand(4, 5, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.rand(8, 5, requires_grad=True, dtype=torch.float64, device=device)
        gradcheck(grad_test_func, inputs=(a, b), check_batched_grad=False)

    def test_views_inherit_ragged_dim(self, device):
        # view
        nt = random_nt_from_dims(
            [4, None, 8, 10], device=device, dtype=torch.float32, layout=torch.jagged
        )
        # inherit ragged dim via -1
        view = nt.view(4, -1, 80)
        self.assertEqual(nt.shape[1], view.shape[1])
        # inherit batch and ragged dims via -1
        view2 = nt.view(-1, -1, 80)
        self.assertEqual(nt.shape[:2], view2.shape[:2])

        # expand
        nt = random_nt_from_dims(
            [3, None, 1], device=device, dtype=torch.float32, layout=torch.jagged
        )
        # inherit batch and ragged dims via -1
        view = nt.expand(-1, -1, 5)
        self.assertEqual(nt.shape[:2], view.shape[:2])

    def test_view_ragged_idx_not_one(self, device):
        nt = random_nt_from_dims(
            [2, None, 20], device=device, dtype=torch.float32, layout=torch.jagged
        )

        view_transposed = nt.transpose(1, 2).view(2, 20, nt.size(1))
        self.assertEqual((2, 20, nt.size(1)), (view_transposed.size()))
        self.assertEqual(view_transposed._base, nt._base)

    def test_unsafe_view(self, device):
        nt = random_nt_from_dims(
            [4, None, 8, 10], device=device, dtype=torch.float32, layout=torch.jagged
        )
        # basic view
        view1 = torch.ops.aten._unsafe_view(nt, (4, -1, 80))
        self.assertEqual((4, nt.size(1), 80), tuple(view1.size()))
        # _unsafe_view differs from view in that the view information is not tracked
        self.assertTrue(view1._base is None)

        # test an unsafe_view when ragged_idx != 1, currently only supports identity view
        nt_t = nt.transpose(1, 2)
        view2 = torch.ops.aten._unsafe_view(nt_t, (4, 8, nt.size(1), 10))
        self.assertEqual((4, 8, nt.size(1), 10), tuple(view2.size()))
        self.assertTrue(view2._base is None)

    @xfailIfTorchDynamo
    @parametrize("requires_grad", [False, True])
    def test_reshape_decomp(self, device, requires_grad):
        # contiguous NT should result in view.
        nt = (
            random_nt_from_dims(
                [3, None, 10],
                device=device,
                dtype=torch.float32,
                layout=torch.jagged,
            )
            .detach()
            .requires_grad_(requires_grad)
        )
        view = nt.reshape(-1, -1, 5, 2)
        self.assertEqual(view.shape[:2], nt.shape[:2])
        self.assertTrue(view._is_view() and view._base is nt)
        # make sure gradients flow back
        if requires_grad:
            view.backward(torch.ones_like(view))
            self.assertEqual(nt.grad, torch.ones_like(nt))

        # non-contiguous NT should result in contiguous copy
        nt = random_nt_from_dims(
            [3, None, 5, 2],
            device=device,
            dtype=torch.float32,
            layout=torch.jagged,
            requires_grad=requires_grad,
        )
        nt_noncontig = nt.transpose(-1, -2)
        self.assertFalse(nt_noncontig.is_contiguous())
        copy = nt_noncontig.reshape(-1, -1, 10)
        self.assertTrue(copy.is_contiguous())
        self.assertEqual(copy.shape[:2], nt.shape[:2])
        # make sure gradients flow back
        if requires_grad:
            copy.backward(torch.ones_like(copy))
            self.assertEqual(nt.grad, torch.ones_like(nt))

    def test_flatten_decomp(self, device):
        nt = random_nt_from_dims(
            [3, None, 5, 2], device=device, dtype=torch.float32, layout=torch.jagged
        )
        flattened = nt.flatten(-2, -1)
        self.assertEqual(flattened.shape, nt.view(3, -1, 10).shape)

        nt = random_nt_from_dims(
            [3, None, 5, 2, 6], device=device, dtype=torch.float32, layout=torch.jagged
        )
        flattened = nt.flatten(-3, -2)
        self.assertEqual(flattened.shape, nt.view(3, -1, 10, 6).shape)

    def test_chunk(self, device):
        # none NJT case
        t = torch.randn(10, 4, 5, requires_grad=True)
        t_list = t.chunk(3, dim=0)
        loss = t_list[0].sum() + t_list[2].sum()
        loss.backward()

        # normal case
        D = 30
        B = 8
        nt = random_nt_from_dims(
            [B, None, D],
            device=device,
            dtype=torch.float32,
            layout=torch.jagged,
            requires_grad=True,
        )
        NUM_CHUNKS = 3
        chunks = nt.chunk(NUM_CHUNKS, dim=-1)
        self.assertEqual(len(chunks), NUM_CHUNKS)
        for i in range(NUM_CHUNKS):
            self.assertEqual(chunks[i].shape[-1], D // NUM_CHUNKS)

        # test chunk_backward
        values = torch.randn(
            5, 11, dtype=torch.float64, device=device, requires_grad=True
        )
        offsets = torch.tensor([0, 2, 3, 5], device=device)

        def grad_test_func(values, offsets):
            nt = torch.nested.nested_tensor_from_jagged(values, offsets)
            chunks = nt.chunk(3, dim=-1)
            return chunks[0].values().sum()

        assert gradcheck(
            grad_test_func,
            inputs=(values, offsets),
            check_batched_grad=False,
        )

        # chunk on batch dim
        chunks = nt.chunk(NUM_CHUNKS, dim=0)
        self.assertEqual(len(chunks), NUM_CHUNKS)
        chunk_size = math.ceil(B / NUM_CHUNKS)
        for i in range(NUM_CHUNKS):
            if i < NUM_CHUNKS - 1:
                self.assertEqual(chunks[i].shape[0], chunk_size)
            else:
                self.assertEqual(chunks[i].shape[0], B - chunk_size * (NUM_CHUNKS - 1))
            offsets_expected = (
                nt._offsets[i * chunk_size + 1 : (i + 1) * chunk_size + 1]
                - nt._offsets[i * chunk_size]
            )
            self.assertEqual(chunks[i]._offsets[1:], offsets_expected)
        self.assertEqual(nt._values, torch.cat([x._values for x in chunks], dim=0))

        with self.assertRaisesRegex(
            RuntimeError,
            "dim != 0 INTERNAL ASSERT FAILED .* Nested Tensor doesn't support chunk backward on dim=0 yet.",
        ):
            # doesn't support backward for chunk (dim=0) yet
            loss = (
                chunks[0].values().sum()
                + chunks[1].values().sum()
                + chunks[2].values().sum()
            )
            loss.backward()

        # chunk on ragged dim not supported
        with self.assertRaisesRegex(
            RuntimeError, "chunk.* not supported for NestedTensor on dim=1"
        ):
            nt.chunk(2, dim=1)

    def test_squeeze(self, device):
        B = 4
        D = 6
        # squeeze middle dim
        nt = random_nt_from_dims(
            [B, None, 1, D], device=device, dtype=torch.float32, layout=torch.jagged
        )
        j0 = nt.shape[1]

        for dim_arg in [-2, 2]:
            out = nt.squeeze(dim_arg)
            self.assertEqual(out.shape, (B, j0, D))
            self.assertEqual(out.unsqueeze(-2), nt)

        # squeeze last dim
        nt = random_nt_from_dims(
            [B, None, 1], device=device, dtype=torch.float32, layout=torch.jagged
        )
        j1 = nt.shape[1]

        for dim_arg in [-1, 2]:
            out = nt.squeeze(dim_arg)
            self.assertEqual(out.shape, (B, j1))
            self.assertEqual(out.unsqueeze(-1), nt)

        # squeeze on batch dim not supported
        with self.assertRaisesRegex(
            RuntimeError, "squeeze.* not supported for NestedTensor on dim=0"
        ):
            nt.squeeze(0)

        # squeeze on ragged dim not supported
        with self.assertRaisesRegex(
            RuntimeError, "squeeze.* not supported for NestedTensor on dim=1"
        ):
            nt.squeeze(1)

    def test_binary_pointwise_broadcasting(self, device):
        # (B, j0, 3, 4)
        ts = self._get_list_for_jagged_tensor(
            ((2, 3, 4), 3, 4), device, requires_grad=True
        )
        # (B, j0, ?, ?) + (?) -> (B, j0, ?, ?)
        # (B, j0, ?, ?) + (?, ?) -> (B, j0, ?, ?)
        # (B, j0, ?, ?) + (1, ?, ?) -> (B, j0, ?, ?)
        # Unsupported: (B, j0, ?, ?) + (1, 1, 1, ?, ?) -> (1, B, j0, ?, ?)
        t_sizes = (
            (4,),
            (1, 4),
            (3, 1),
            (1, 3, 1),
            (1, 1, 1, 4),
            # (1, 1, 1, 1, 4), (unsupported today)
        )

        def grad_test_func(t, *ts):
            nt = torch.nested.as_nested_tensor(list(ts), layout=torch.jagged)
            out = nt + t
            return out.values()

        for t_size in t_sizes:
            t = torch.rand(
                t_size, requires_grad=True, device=device, dtype=torch.float64
            )
            gradcheck(grad_test_func, inputs=(t, *ts), check_batched_grad=False)

    def test_threshold_backward(self, device):
        ts1 = self._get_list_for_jagged_tensor(
            ((2, 3, 4), 16), device=device, requires_grad=False
        )
        ts2 = self._get_list_for_jagged_tensor(
            ((2, 3, 4), 16), device=device, requires_grad=False
        )

        nt1, offsets = jagged_from_list(ts1, None)
        nt2, offsets = jagged_from_list(ts2, offsets)
        buf1 = nt1.values().detach().clone()
        buf2 = nt2.values().detach().clone()

        res_nt = torch.ops.aten.threshold_backward(nt1, nt2, 0.0)
        res_dense = torch.ops.aten.threshold_backward(buf1, buf2, 0.0)

        self.assertEqual(res_dense, res_nt.values())

    @dtypes(torch.float32)
    @parametrize(
        "func",
        [torch.ops.aten.sum.dim_IntList, torch.ops.aten.mean.dim],
        name_fn=get_op_name,
    )
    @parametrize("keepdim", [False, True])
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_jagged_op_different_output_shape_dim(
        self, device, dtype, keepdim, requires_grad, components_require_grad, func
    ):
        """
        Operator passes when reducing on valid reduction dimensions.
        This test is for operators which return an output tensor with a shape different from the input tensor.
        """
        if get_op_name(func) == "mean" and not keepdim:
            return

        op_name = get_op_name(func)

        ts = self._get_list_for_jagged_tensor(
            ((2, 3, 4), 3, 4), device=device, requires_grad=True
        )  # (B, j0, 3, 4)

        # verify correctness of shapes (assuming that ragged_idx == 1)
        if op_name == "sum":
            reduce_dims = (
                ((0, 1), (3, 4), (1, 1, 3, 4), (0,)),  # batch, ragged
                ((2, 3), (3, None), (3, None, 1, 1), (1, 2)),  # non-batch, non-batch
                ((0, 1, 3), (3,), (1, 1, 3, 1), (0, 2)),  # batch, ragged, non-batch
                ((0, 1, 2), (4,), (1, 1, 1, 4), (0, 1)),  # batch, ragged, non-batch
                (
                    (0, 1, 2, 3),
                    (),
                    (1, 1, 1, 1),
                    (0, 1, 2),
                ),  # batch, ragged, non-batch, non-batch
                ((2,), (3, None, 4), (3, None, 1, 4), (1,)),  # non-batch
            )  # (dims, expected shape, expected keepdim shape, reduce_dim_expected), where j0 is represented as None
        elif op_name == "mean":
            reduce_dims = (
                ((2,), (3, None, 4), (3, None, 1, 4), (1,)),
                ((3,), (3, None, 3), (3, None, 3, 1), (2,)),
            )

        for rd, ref_shape_no_keepdim, ref_shape_keepdim, _ in reduce_dims:
            nt = torch.nested.as_nested_tensor(ts, layout=torch.jagged)
            out = func(nt, dim=rd, keepdim=keepdim)
            ref_shape = ref_shape_keepdim if keepdim else ref_shape_no_keepdim
            if not torch.compiler.is_compiling:  # if not using torch dynamo
                self.assertEqual(len(out.shape), len(ref_shape))
                for o, r in zip(out.shape, ref_shape):
                    if r is not None:
                        self.assertEqual(o, r)
                    else:
                        self.assertTrue(isinstance(o, torch.SymInt))

        # verify correctness of values
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,
        )
        for tensor_list, reduce_dim_tuple in itertools.product(
            tensor_lists, reduce_dims
        ):
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            reduce_dim, _, _, reduce_dim_expected = reduce_dim_tuple

            if nt.dim() > reduce_dim[-1]:
                out_actual = func(nt, dim=reduce_dim, keepdim=keepdim)
                if nt._ragged_idx in reduce_dim:  # raggedness reduced away
                    out_expected = func(
                        nt.values(), dim=reduce_dim_expected, keepdim=keepdim
                    )
                    self.assertTrue(torch.allclose(out_actual, out_expected))
                else:  # raggedness preserved
                    out_expected = func(nt.values(), dim=reduce_dim_expected)
                    self.assertTrue(
                        torch.allclose(
                            out_actual.values().view(-1), out_expected.view(-1)
                        )
                    )

    @dtypes(torch.float32)
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_softmax_dim(
        self,
        device,
        dtype,
        requires_grad,
        components_require_grad,
    ):
        """
        Softmax passes when reducing on valid reduction dimensions.
        """
        ts = self._get_list_for_jagged_tensor(
            ((2, 3, 4), 3, 4), device=device, requires_grad=True
        )  # (B, j0, 3, 4)

        output_shape = (3, None, 3, 4)

        # verify correctness of shapes (assuming that ragged_idx == 1)
        reduce_dims = (
            (2, 1),
            (3, 2),
        )  # (reduction dimension, effective reduction dimension for baseline)

        for reduce_dim, _ in reduce_dims:
            nt = torch.nested.as_nested_tensor(ts, layout=torch.jagged)
            out_actual = torch.nn.functional.softmax(nt, dim=reduce_dim)
            torch._dynamo.disable(self.assertEqual)(
                len(out_actual.shape), len(output_shape)
            )  # disable if running on dynamo
            for dim_actual, dim_expected in zip(out_actual.shape, output_shape):
                if dim_expected is not None:
                    self.assertEqual(dim_actual, dim_expected)
                else:
                    self.assertTrue(isinstance(dim_actual, torch.SymInt))

        # verify correctness of values
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,
        )
        for tensor_list, reduce_dim_tuple in itertools.product(
            tensor_lists, reduce_dims
        ):
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            reduce_dim, reduce_dim_expected = reduce_dim_tuple

            if nt.dim() > reduce_dim:
                out_actual = torch.nn.functional.softmax(
                    nt, dim=reduce_dim
                )  # nested tensor
                out_expected = torch.nn.functional.softmax(
                    nt.values(), dim=reduce_dim_expected
                )  # dense tensor of dimensions 1 less than out_actual
                self.assertTrue(
                    torch.allclose(out_actual.values().view(-1), out_expected.view(-1))
                )

    @dtypes(torch.float32)
    @parametrize(
        "func",
        [torch.ops.aten.sum.dim_IntList, torch.ops.aten.mean.dim],
        name_fn=get_op_name,
    )
    @parametrize("keepdim", [False, True])
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_op_dim_reduce_ragged_idx_1_different_output_shape(
        self, device, dtype, keepdim, requires_grad, components_require_grad, func
    ):
        """
        Operator on NestedTensor passes when trying to reduce across ragged dimension, where ragged_idx == 1.
        This test is for operators which return an output tensor with a shape different from the input tensor.
        """
        if get_op_name(func) == "mean" and not keepdim:
            return

        op_name = get_op_name(func)

        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,  # (B, *, 1)
        )
        reduce_dim = (1,)  # ragged

        for tensor_list in tensor_lists:
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            out_actual = func(nt, dim=reduce_dim, keepdim=keepdim)
            out_expected = torch.cat(
                [func(t, dim=(reduce_dim[0] - 1)).unsqueeze(0) for t in nt.unbind()]
            )

            self.assertFalse(
                out_actual.is_nested,
                f"{op_name}(): the result of reducing a nested tensor along the ragged dimension is a dense tensor",
            )  # output is a dense tensor
            self.assertTrue(torch.allclose(out_actual, out_expected))

    @dtypes(torch.float32)
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_softmax_dim_reduce_ragged_idx_1(
        self, device, dtype, requires_grad, components_require_grad
    ):
        """
        Softmax on NestedTensor passes when trying to reduce across ragged dimension, where ragged_idx == 1.
        """
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,  # (B, *, 1)
            include_2d_tensor=True,  # (B, *)
        )
        reduce_dim = 1  # ragged

        for tensor_list in tensor_lists:
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            out_actual = torch.nn.functional.softmax(nt, dim=reduce_dim)
            out_expected = torch.cat(
                [
                    torch.nn.functional.softmax(t, dim=reduce_dim - 1)
                    for t in nt.unbind()
                ]
            )

            self.assertTrue(
                out_actual.is_nested,
                "softmax(): the result of reducing a nested tensor along the ragged dimension is a nested tensor",
            )  # output is a nested tensor
            self.assertTrue(torch.allclose(out_actual.values(), out_expected))

    @dtypes(torch.float32)
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_softmax_reduce_batch_dim(
        self, device, dtype, requires_grad, components_require_grad
    ):
        """
        Softmax on NestedTensor fails when trying to reduce across batch dimension.
        """
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,  # (B, *, 1)
        )
        reduce_dim = 0  # batch

        for tensor_list in tensor_lists:
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            with self.assertRaisesRegex(
                RuntimeError,
                "not supported when reducing across the batch dimension for NestedTensor",
            ):
                out = torch.nn.functional.softmax(nt, dim=reduce_dim)

    @dtypes(torch.float32)
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_layer_norm_reduce_ragged_idx_1(
        self, device, dtype, requires_grad, components_require_grad
    ):
        """
        Layer normalization on NestedTensor passes when trying to normalize across ragged dimension, where ragged_idx == 1.
        """

        # requires_grad = False does not currently work with dynamo tests and throws this error:
        #   AssertionError: SymInts must use SymNodeVariable.
        #   If the underlying value is static, we will create a ConstantVariable and specialize.
        if torch._dynamo.is_compiling() and not requires_grad:
            return

        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,  # (B, *, 1)
        )

        for tensor_list in tensor_lists:
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            if (
                nt.dim() >= 3
            ):  # layer norm only works for tensors with 3 or more dimensions
                normalized_shape = nt.shape[nt._ragged_idx :]

                out_actual = torch.nn.functional.layer_norm(
                    nt, normalized_shape=normalized_shape
                )
                out_expected = torch.cat(
                    [
                        torch.nn.functional.layer_norm(t, normalized_shape=t.shape)
                        for t in nt.unbind()
                    ]
                )  # e.g. in 3D tensor (B, *, M), performs layer normalization on B 2D tensors (*, M)

                self.assertTrue(
                    out_actual.is_nested,
                    "layer_norm(): the result of reducing a nested tensor along the ragged dimension is a nested tensor",
                )  # output is a nested tensor
                self.assertEqual(out_actual._values.shape, out_expected.shape)
                self.assertTrue(torch.allclose(out_actual.values(), out_expected))

    @dtypes(torch.float32)
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_layer_norm_2d_input(
        self,
        device,
        dtype,
        requires_grad,
        components_require_grad,
    ):
        """
        Layer normalization on NestedTensor fails when trying to operate on a 2-dimensional tensor
        """
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,  # (B, *, 1)
            include_2d_tensor=True,  # (B, *)
        )

        for tensor_list in tensor_lists:
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            if nt.dim() <= 2:
                with self.assertRaisesRegex(
                    RuntimeError,
                    "not supported for NestedTensor objects with 2 or fewer dimensions",
                ):
                    out = torch.nn.functional.layer_norm(
                        nt, normalized_shape=(nt.shape[nt._ragged_idx],)
                    )

    @dtypes(torch.float32)
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_layer_norm_operate_on_batch_dim(
        self,
        device,
        dtype,
        requires_grad,
        components_require_grad,
    ):
        """
        Layer normalization on NestedTensor fails when trying to operate on the batch dimension
        """
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,  # (B, *, 1)
            include_2d_tensor=True,  # (B, *)
        )

        for tensor_list in tensor_lists:
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            if nt.dim() > 2:  # cannot perform layer normalization on 2D tensors
                with self.assertRaisesRegex(
                    RuntimeError,
                    "not supported when normalizing over the batch dimension for NestedTensor",
                ):
                    out = torch.nn.functional.layer_norm(nt, normalized_shape=nt.shape)

    @dtypes(torch.float32)
    @parametrize(
        "func",
        [torch.ops.aten.sum.dim_IntList, torch.ops.aten.mean.dim],
        name_fn=get_op_name,
    )
    @parametrize(
        "transpose_offset", [1, 2]
    )  # [transpose consecutive dimensions, transpose nonconsecutive dimensions]
    @parametrize("keepdim", [False, True])
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_op_dim_reduce_ragged_idx_greater_than_1_different_output_shape(
        self,
        device,
        dtype,
        keepdim,
        requires_grad,
        components_require_grad,
        func,
        transpose_offset,
    ):
        """
        Operator on NestedTensor passes when trying to reduce across a transposed ragged dimension, i.e. ragged_idx > 1
        This test is for operators which return an output tensor with a shape different from the input tensor.
        """
        if get_op_name(func) == "mean" and not keepdim:
            return

        op_name = get_op_name(func)

        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,  # (B, *, 1)
            include_2d_tensor=True,  # (B, *)
        )

        for tensor_list in tensor_lists:
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            if nt.dim() > nt._ragged_idx + transpose_offset:
                nt_transposed = nt.transpose(
                    nt._ragged_idx, nt._ragged_idx + transpose_offset
                )
                reduce_dim = (nt_transposed._ragged_idx,)  # ragged

                out_actual = func(nt_transposed, dim=reduce_dim, keepdim=keepdim)
                out_expected = torch.cat(
                    [
                        func(t, dim=(reduce_dim[0] - 1)).unsqueeze(0)
                        for t in nt_transposed.unbind()
                    ]
                )

                self.assertFalse(
                    out_actual.is_nested,
                    f"{op_name}(): the result of reducing a nested tensor along the ragged dimension is a dense tensor",
                )  # output is a dense tensor
                self.assertTrue(torch.allclose(out_actual, out_expected, rtol=1e-4))

    @dtypes(torch.float32)
    @parametrize(
        "transpose_offset", [1, 2]
    )  # [transpose consecutive dimensions, transpose nonconsecutive dimensions]
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_softmax_dim_reduce_ragged_idx_greater_than_1_same_output_shape(
        self,
        device,
        dtype,
        requires_grad,
        components_require_grad,
        transpose_offset,
    ):
        """
        Softmax on NestedTensor fails when trying to reduce across a transposed ragged dimension, i.e. ragged_idx > 1
        This test is for operators which return an output tensor with the same shape as the input tensor.
        """
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,  # (B, *, 1)
        )

        for tensor_list in tensor_lists:
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            if nt.dim() > nt._ragged_idx + transpose_offset:
                nt_transposed = nt.transpose(
                    nt._ragged_idx, nt._ragged_idx + transpose_offset
                )
                reduce_dim = nt_transposed._ragged_idx  # ragged

                with self.assertRaisesRegex(
                    RuntimeError,
                    "not supported when reducing along the ragged dimension for ragged_idx > 1 for NestedTensor",
                ):
                    out = torch.nn.functional.softmax(nt_transposed, dim=reduce_dim)

    @dtypes(torch.float32)
    @parametrize(
        "func",
        [torch.ops.aten.sum.dim_IntList, torch.ops.aten.mean.dim],
        name_fn=get_op_name,
    )
    @parametrize("keepdim", [False, True])
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_op_dim_transpose_non_ragged_dim_different_output_shape(
        self, device, dtype, keepdim, requires_grad, components_require_grad, func
    ):
        """
        Operator passes when reducing transposed nested tensors on valid reduction dimensions.
        This test is for operators which return an output tensor with a shape different from the input tensor.
        """
        if get_op_name(func) == "mean" and not keepdim:
            return

        # verify correctness of shapes (assuming that ragged_idx == 1)
        if get_op_name(func) == "sum":
            reduce_dims = (
                ((0, 1), (3, 4), (1, 1, 3, 4), (0,)),  # batch, ragged
                ((2, 3), (3, None), (3, None, 1, 1), (1, 2)),  # non-batch, non-batch
                ((0, 1, 3), (3,), (1, 1, 3, 1), (0, 2)),  # batch, ragged, non-batch
                ((0, 1, 2), (4,), (1, 1, 1, 4), (0, 1)),  # batch, ragged, non-batch
                (
                    (0, 1, 2, 3),
                    (),
                    (1, 1, 1, 1),
                    (0, 1, 2),
                ),  # batch, ragged, non-batch, non-batch
                ((2,), (3, None, 4), (3, None, 1, 4), (1,)),  # non-batch
            )  # (dims, expected shape, expected keepdim shape, reduce_dim_expected), where j0 is represented as None
        elif get_op_name(func) == "mean":
            reduce_dims = (
                ((2,), (3, None, 4), (3, None, 1, 4), (1,)),
                ((3,), (3, None, 3), (3, None, 3, 1), (2,)),
            )

        # verify correctness of values
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
        )
        for tensor_list, reduce_dim_tuple in itertools.product(
            tensor_lists, reduce_dims
        ):
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            ).transpose(-1, -2)

            reduce_dim, _, _, reduce_dim_expected = reduce_dim_tuple

            if nt.dim() > max(
                reduce_dim[-1], nt._ragged_idx + 2
            ):  # ensure that transposed dimensions are non-batch, non-ragged dimensions
                out_actual = func(nt, dim=reduce_dim, keepdim=keepdim)
                if nt._ragged_idx in reduce_dim:  # raggedness reduced away
                    out_expected = func(
                        nt.values(), dim=reduce_dim_expected, keepdim=keepdim
                    )
                    self.assertTrue(torch.allclose(out_actual, out_expected))
                else:  # raggedness preserved
                    out_expected = func(nt.values(), dim=reduce_dim_expected)
                    self.assertTrue(
                        torch.allclose(
                            out_actual.values().view(-1), out_expected.view(-1)
                        )
                    )

    @dtypes(torch.float32)
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_softmax_dim_transpose_non_ragged_dim(
        self,
        device,
        dtype,
        requires_grad,
        components_require_grad,
    ):
        """
        Softmax passes when reducing transposed nested tensors on valid reduction dimensions.
        This test is for operators which return an output tensor with the same shape as the input tensor.
        """
        # verify correctness of shapes (assuming that ragged_idx == 1)
        reduce_dims = (
            (2, 1),
            (3, 2),
        )  # (reduction dimension, effective reduction dimension for baseline)

        # verify correctness of values
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False,
            include_requires_grad=components_require_grad,
            include_inner_dim_size_1=True,  # (B, *, 1)
        )
        for tensor_list, reduce_dim_tuple in itertools.product(
            tensor_lists, reduce_dims
        ):
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            ).transpose(-1, -2)

            reduce_dim, reduce_dim_expected = reduce_dim_tuple

            if nt.dim() > max(reduce_dim, nt._ragged_idx + 2):
                out_actual = torch.nn.functional.softmax(
                    nt, dim=reduce_dim
                )  # nested tensor
                out_expected = torch.nn.functional.softmax(
                    nt.values(), dim=reduce_dim_expected
                )  # dense tensor of dimensions 1 less than out_actual

                self.assertTrue(
                    torch.allclose(out_actual.values().view(-1), out_expected.view(-1))
                )

    @dtypes(torch.float32)
    @parametrize("keepdim", [False, True])
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_sum_dim_reduce_ragged_and_non_batch(
        self,
        device,
        dtype,
        keepdim,
        requires_grad,
        components_require_grad,
    ):
        """
        Sum on NestedTensor fails when trying to reduce across ragged and non-batch dimensions
        """
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False, include_requires_grad=components_require_grad
        )
        reduce_dims = (
            (1, 2),  # ragged, non-batch
            (1, 3),  # ragged, non-batch
        )

        for tensor_list, reduce_dim in itertools.product(tensor_lists, reduce_dims):
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            if nt.dim() > reduce_dim[-1]:
                with self.assertRaisesRegex(
                    RuntimeError,
                    "not supported along a ragged and non-batch dimension for NestedTensor",
                ):
                    out = torch.sum(nt, dim=reduce_dim, keepdim=keepdim)

    @dtypes(torch.float32)
    @parametrize("keepdim", [False, True])
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_sum_dim_reduce_batch_and_non_batch(
        self,
        device,
        dtype,
        keepdim,
        requires_grad,
        components_require_grad,
    ):
        """
        Sum on NestedTensor fails when trying to reduce across batch and non-batch dimensions
        """
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False, include_requires_grad=components_require_grad
        )
        reduce_dims = (
            (0, 2),  # batch, non-batch
            (0, 3),  # batch, non-batch
        )

        for tensor_list, reduce_dim in itertools.product(tensor_lists, reduce_dims):
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            if nt.dim() > reduce_dim[-1]:
                with self.assertRaisesRegex(
                    RuntimeError,
                    "not supported along the batch dimension but not the ragged dimension for NestedTensor",
                ):
                    out = torch.sum(nt, dim=reduce_dim, keepdim=keepdim)

    @dtypes(torch.float32)
    @parametrize(
        "func",
        [torch.ops.aten.sum.dim_IntList, torch.ops.aten.mean.dim],
        name_fn=get_op_name,
    )
    @parametrize("keepdim", [False, True])
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_op_dim_reduce_batch_only_different_output_shape(
        self, device, dtype, keepdim, requires_grad, components_require_grad, func
    ):
        """
        Operator on NestedTensor fails when trying to reduce across batch dimension
        """
        if get_op_name(func) == "mean" and not keepdim:
            return

        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False, include_requires_grad=components_require_grad
        )
        reduce_dim = (0,)  # batch

        for tensor_list in tensor_lists:
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            with self.assertRaisesRegex(
                RuntimeError,
                "not supported along the batch dimension but not the ragged dimension for NestedTensor",
            ):
                out = func(nt, dim=reduce_dim, keepdim=keepdim)

    @dtypes(torch.float32)
    @parametrize(
        "func",
        [torch.ops.aten.sum.dim_IntList, torch.ops.aten.mean.dim],
        name_fn=get_op_name,
    )
    @parametrize("keepdim", [False, True])
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_op_dim_with_lengths_different_output_shape(
        self,
        device,
        dtype,
        keepdim,
        requires_grad,
        components_require_grad,
        func,
    ):
        """
        Operator on NestedTensor fails when trying to reduce a nested tensor with lengths,
        i.e. a nested tensor with holes, if reducing on the ragged dimension.
        This test is for operators which return an output tensor with different shape than the input tensor.
        """
        if get_op_name(func) == "mean" and not keepdim:
            return

        reduce_dims = ((1,), (2,), (2, 3))

        lengths = torch.randint(5, 10, (20,), device=device)
        offsets = torch.zeros((21,), device=device, dtype=torch.int)
        torch.cumsum(lengths, dim=0, out=offsets[1:])

        values = torch.randn(
            (offsets[-1].item(), 20),
            device=device,
            dtype=dtype,
            requires_grad=requires_grad,
        )

        nt_with_holes = torch.nested.nested_tensor_from_jagged(
            values,
            offsets,
            lengths=offsets.diff() - 2,  # arbitrary subtraction to create holes
        )

        for reduce_dim in reduce_dims:
            if nt_with_holes.dim() > reduce_dim[-1]:
                if nt_with_holes._ragged_idx in reduce_dim:
                    with self.assertRaisesRegex(
                        RuntimeError,
                        "not supported where lengths is not None "
                        + "if reducing across the ragged dimension for NestedTensor",
                    ):
                        out = func(nt_with_holes, dim=reduce_dim, keepdim=keepdim)
                else:
                    out = func(nt_with_holes, dim=reduce_dim, keepdim=keepdim)

    @dtypes(torch.float32)
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_softmax_dim_with_lengths(
        self,
        device,
        dtype,
        requires_grad,
        components_require_grad,
    ):
        """
        Softmax on NestedTensor fails when trying to reduce a nested tensor with lengths,
        i.e. a nested tensor with holes, if reducing on the ragged dimension.
        """
        reduce_dims = (1, 2, 3)

        lengths = torch.randint(5, 10, (20,), device=device)
        offsets = torch.zeros((21,), device=device, dtype=torch.int)
        torch.cumsum(lengths, dim=0, out=offsets[1:])

        values = torch.randn(
            (offsets[-1].item(), 20),
            device=device,
            dtype=dtype,
            requires_grad=requires_grad,
        )

        nt_with_holes = torch.nested.nested_tensor_from_jagged(
            values,
            offsets,
            lengths=offsets.diff() - 2,  # arbitrary subtraction to create holes
        )

        for reduce_dim in reduce_dims:
            if nt_with_holes.dim() > reduce_dim:
                if nt_with_holes._ragged_idx == reduce_dim:
                    with self.assertRaisesRegex(
                        RuntimeError,
                        "not supported where lengths is not None "
                        + "if reducing across the ragged dimension for NestedTensor",
                    ):
                        out = torch.nn.functional.softmax(nt_with_holes, dim=reduce_dim)
                else:
                    out = torch.nn.functional.softmax(nt_with_holes, dim=reduce_dim)

    @skipIfTorchDynamo(
        "ragged_size = nt_with_holes.shape[nt_with_holes._ragged_idx] does not currently work "
        + "with dynamo tests and throws this error: `AssertionError: SymInts must use SymNodeVariable. "
        + "If the underlying value is static, we will create a ConstantVariable and specialize.`"
    )
    @dtypes(torch.float32)
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_layer_norm_with_lengths(
        self,
        device,
        dtype,
        requires_grad,
        components_require_grad,
    ):
        """
        Layer normalization on NestedTensor fails when trying to operate on a nested tensor with lengths,
        i.e. a nested tensor with holes, if operating on the ragged dimension.
        """

        # create components for nested tensor
        lengths = torch.randint(5, 10, (20,), device=device)
        offsets = torch.zeros((21,), device=device, dtype=torch.int)
        torch.cumsum(lengths, dim=0, out=offsets[1:])
        values = torch.randn(
            (offsets[-1].item(), 10, 30),
            device=device,
            dtype=dtype,
            requires_grad=requires_grad,
        )

        nt_with_holes = torch.nested.nested_tensor_from_jagged(
            values,
            offsets,
            lengths=offsets.diff() - 2,  # arbitrary subtraction to create holes
        )

        ragged_size = nt_with_holes.shape[nt_with_holes._ragged_idx]

        normalized_shapes = (
            (10, 30),  # normalization on non-ragged dimension passes
            (ragged_size, 10, 30),  # normalization on ragged dimension fails
        )

        for normalized_shape in normalized_shapes:
            if ragged_size in normalized_shape:
                with self.assertRaisesRegex(
                    RuntimeError,
                    "not supported where lengths is not None if operating on the ragged dimension for NestedTensor",
                ):
                    out = torch.nn.functional.layer_norm(
                        nt_with_holes, normalized_shape=normalized_shape
                    )
            else:
                out = torch.nn.functional.layer_norm(
                    nt_with_holes, normalized_shape=normalized_shape
                )

    @dtypes(torch.float32)
    @parametrize("keepdim", [True])
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_mean_dim_reduce_multiple_dims(
        self,
        device,
        dtype,
        keepdim,
        requires_grad,
        components_require_grad,
    ):
        """
        Mean on NestedTensor fails when trying to reduce across multiple dimensions
        """
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False, include_requires_grad=components_require_grad
        )
        reduce_dims = ((0, 1), (2, 3), (2, 3, 4))

        for tensor_list, reduce_dim in itertools.product(tensor_lists, reduce_dims):
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            if nt.dim() > reduce_dim[-1]:
                with self.assertRaisesRegex(
                    RuntimeError,
                    "not supported across multiple dimensions for NestedTensor",
                ):
                    out = torch.mean(nt, dim=reduce_dim, keepdim=keepdim)

    @dtypes(torch.float32)
    @parametrize("keepdim", [False, True])
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_mean_dim_keepdim_False(
        self,
        device,
        dtype,
        keepdim,
        requires_grad,
        components_require_grad,
    ):
        """
        Mean on NestedTensor fails when keepdim=False
        """
        tensor_lists = self._get_example_tensor_lists(
            include_list_of_lists=False, include_requires_grad=components_require_grad
        )
        reduce_dims = ((1,), (2,), (3,))

        for tensor_list, reduce_dim in itertools.product(tensor_lists, reduce_dims):
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            if nt.dim() > reduce_dim[-1]:
                if not keepdim:
                    with self.assertRaisesRegex(
                        RuntimeError,
                        "not supported when keepdim=False for NestedTensor",
                    ):
                        out = torch.mean(nt, dim=reduce_dim, keepdim=keepdim)
                else:
                    out = torch.mean(nt, dim=reduce_dim, keepdim=keepdim)

    @dtypes(torch.float, torch.double, torch.half)
    @parametrize("requires_grad", [False, True])
    @parametrize("weights_only", [False, True])
    def test_serialization(self, device, dtype, requires_grad, weights_only):
        def compare_metadata(nt1, nt2):
            self.assertEqual(nt1._nested_tensor_size(), nt2._nested_tensor_size())
            self.assertEqual(nt1._nested_tensor_strides(), nt2._nested_tensor_strides())
            self.assertEqual(
                nt1._nested_tensor_storage_offsets(),
                nt2._nested_tensor_storage_offsets(),
            )

        nt_contiguous, nt_noncontiguous = random_nt_noncontiguous_pair((2, 3, 6, 7))
        for a in [nt_contiguous, nt_noncontiguous]:
            buffer = io.BytesIO()
            serialized = torch.save(a, buffer)
            buffer.seek(0)
            b = torch.load(buffer, weights_only=weights_only)
            # should be both conceptually equal and metadata equivalent
            self.assertEqual(a, b)
            compare_metadata(a, b)
            # should be conceptually equal but not necessarily metadata equivalent
            self.assertEqual(b, nt_contiguous)
            self.assertEqual(b, nt_noncontiguous)

    @unittest.skipIf(
        PYTORCH_CUDA_MEMCHECK, "is_pinned uses failure to detect pointer property"
    )
    @onlyCUDA
    def test_pin_memory(self, device):
        nt_contiguous, nt_noncontiguous = random_nt_noncontiguous_pair((2, 3, 6, 7))
        for nt in [nt_contiguous, nt_noncontiguous]:
            self.assertFalse(nt.is_pinned())
            pinned = nt.pin_memory(device)
            self.assertTrue(pinned.is_pinned())
            self.assertEqual(nt, pinned)
            self.assertNotEqual(nt.data_ptr(), pinned.data_ptr())
            # test that pin_memory on already pinned tensor has no effect
            self.assertIs(pinned, pinned.pin_memory())
            self.assertEqual(pinned.data_ptr(), pinned.pin_memory().data_ptr())

    @torch.compiler.disable
    def _validate_nt(
        self,
        nt,
        device,
        dtype,
        layout,
        requires_grad,
        dim,
        batch_size,
        contiguous,
        cached_min_seqlen=None,
        cached_max_seqlen=None,
        base=None,
        ref_nt=None,
    ):
        # Validate a bunch of properties after NT construction.
        device = torch.device(device)
        self.assertEqual(nt.dim(), dim)
        self.assertEqual(nt.device, device)
        self.assertEqual(nt.dtype, dtype)
        self.assertEqual(nt.layout, layout)
        self.assertEqual(nt.requires_grad, requires_grad)
        self.assertEqual(nt.is_contiguous(), contiguous)

        if layout == torch.jagged:
            self.assertEqual(nt._values.device, device)
            self.assertEqual(nt._offsets.device, device)
            self.assertEqual(nt.shape[0], batch_size)
            self.assertTrue(isinstance(nt.shape[1], torch.SymInt))

            if base is not None:
                self.assertTrue(nt._is_view() and nt._base is base)
                replay_cache = nt._view_func(torch.randn_like(nt._base))._metadata_cache
                self.assertEqual(
                    "min_seqlen" in replay_cache, cached_min_seqlen is not None
                )
                self.assertEqual(
                    "max_seqlen" in replay_cache, cached_max_seqlen is not None
                )

            self.assertEqual(
                "min_seqlen" in nt._metadata_cache, cached_min_seqlen is not None
            )
            self.assertEqual(
                "max_seqlen" in nt._metadata_cache, cached_max_seqlen is not None
            )

            if cached_min_seqlen is not None:
                self.assertEqual(nt._min_seqlen, cached_min_seqlen)

            if cached_max_seqlen is not None:
                self.assertEqual(nt._max_seqlen, cached_max_seqlen)

        if ref_nt is not None:
            self.assertEqual(nt.size(0), ref_nt.size(0))
            for n1, n2 in zip(nt.unbind(), ref_nt.unbind()):
                self.assertEqual(n1, n2)

    @dtypes(torch.float, torch.double, torch.half)
    @parametrize("requires_grad", [False, True])
    @parametrize("components_require_grad", [False, True])
    def test_jagged_layout_construction_nested_tensor(
        self, device, dtype, requires_grad, components_require_grad
    ):
        for tensor_list in self._get_example_tensor_lists(
            include_list_of_lists=True, include_requires_grad=components_require_grad
        ):
            nt = torch.nested.nested_tensor(
                tensor_list,
                device=device,
                dtype=dtype,
                layout=torch.jagged,
                requires_grad=requires_grad,
            )

            expected_dim = torch.as_tensor(tensor_list[0]).dim() + 1
            expected_batch_size = len(tensor_list)
            expected_contiguous = True
            expected_min_seqlen = min(
                (torch.tensor(t) if isinstance(t, list) else t).shape[0]
                for t in tensor_list
            )
            expected_max_seqlen = max(
                (torch.tensor(t) if isinstance(t, list) else t).shape[0]
                for t in tensor_list
            )
            self._validate_nt(
                nt,
                device,
                dtype,
                torch.jagged,
                requires_grad,
                expected_dim,
                expected_batch_size,
                expected_contiguous,
                expected_min_seqlen,
                expected_max_seqlen,
            )

            # Make sure grads -don't- flow back into original tensors for nested_tensor()
            if requires_grad:
                (nt * 2).backward(torch.ones_like(nt))
            for t in tensor_list:
                t = t if isinstance(t, torch.Tensor) else torch.as_tensor(t)
                self.assertTrue(t.grad is None)

    @dtypes(torch.float, torch.double, torch.half)
    @parametrize("components_require_grad", [False, True])
    def test_jagged_layout_construction_as_nested_tensor(
        self, device, dtype, components_require_grad
    ):
        # NB: as_nested_tensor(tensor_list) doesn't support lists of lists for tensor_list
        for tensor_list in self._get_example_tensor_lists(
            include_list_of_lists=False, include_requires_grad=components_require_grad
        ):
            nt = torch.nested.as_nested_tensor(
                tensor_list, device=device, dtype=dtype, layout=torch.jagged
            )

            # nt.requires_grad=True should be set if at least one component requires grad
            expected_dim = tensor_list[0].dim() + 1
            expected_batch_size = len(tensor_list)
            expected_contiguous = True
            expected_min_seqlen = min(
                (torch.tensor(t) if isinstance(t, list) else t).shape[0]
                for t in tensor_list
            )
            expected_max_seqlen = max(
                (torch.tensor(t) if isinstance(t, list) else t).shape[0]
                for t in tensor_list
            )
            self._validate_nt(
                nt,
                device,
                dtype,
                torch.jagged,
                components_require_grad,
                expected_dim,
                expected_batch_size,
                expected_contiguous,
                expected_min_seqlen,
                expected_max_seqlen,
            )

            # Make sure grads flow back into original tensors for as_nested_tensor()
            if components_require_grad:
                (nt * 2).backward(torch.ones_like(nt))
                for t in tensor_list:
                    if t.requires_grad:
                        self.assertEqual(t.grad, torch.ones_like(t) * 2)
                    else:
                        self.assertTrue(t.grad is None)

    @xfailIfTorchDynamo
    @unittest.skipIf(
        PYTORCH_CUDA_MEMCHECK, "is_pinned uses failure to detect pointer property"
    )
    @onlyCUDA
    def test_jagged_layout_construction_with_pinned_memory(self, device):
        for tensor_list in self._get_example_tensor_lists():
            nt = torch.nested.nested_tensor(
                tensor_list, layout=torch.jagged, device="cpu", pin_memory=True
            )

            expected_dim = torch.as_tensor(tensor_list[0]).dim() + 1
            expected_batch_size = len(tensor_list)
            expected_min_seqlen = min(
                (torch.tensor(t) if isinstance(t, list) else t).shape[0]
                for t in tensor_list
            )
            expected_max_seqlen = max(
                (torch.tensor(t) if isinstance(t, list) else t).shape[0]
                for t in tensor_list
            )
            self._validate_nt(
                nt,
                device="cpu",
                dtype=torch.float32,
                layout=torch.jagged,
                requires_grad=False,
                dim=expected_dim,
                batch_size=expected_batch_size,
                contiguous=True,
                cached_min_seqlen=expected_min_seqlen,
                cached_max_seqlen=expected_max_seqlen,
            )
            self.assertTrue(nt.is_pinned())

    @dtypes(torch.float, torch.double, torch.half)
    @parametrize("requires_grad", [False, True])
    @parametrize("values_is_view", [False, True])
    def test_jagged_view_from_values_offsets(
        self, device, dtype, requires_grad, values_is_view
    ):
        if values_is_view:
            # make values a view of base
            base = torch.randn(
                2, 3, 4, 5, 6, device=device, dtype=dtype, requires_grad=requires_grad
            )
            values = base.flatten(0, -2)
        else:
            values = torch.randn(
                10, 5, device=device, dtype=dtype, requires_grad=requires_grad
            )
        offsets = torch.tensor([0, 2, 4, 6, 10], device=device, dtype=torch.int64)

        nt = nested_view_from_values_offsets(values, offsets)

        expected_dim = values.dim() + 1
        expected_batch_size = offsets.shape[0] - 1
        expected_base = base if values_is_view else values
        lengths = offsets.diff()
        self._validate_nt(
            nt,
            device,
            dtype,
            torch.jagged,
            requires_grad,
            expected_dim,
            expected_batch_size,
            # ensure NT is a proper view
            base=expected_base,
            contiguous=True,
            # if no min / max are passed, expect the metadata cache to be empty
            cached_min_seqlen=None,
            cached_max_seqlen=None,
        )

        if requires_grad:
            # Make sure grads flow back
            (nt * 2).backward(torch.ones_like(nt))

            @torch.compiler.disable
            def _check_grad(t):
                self.assertTrue(t.grad is not None)
                self.assertEqual(t.grad, torch.ones_like(t) * 2)

            _check_grad(base if values_is_view else values)

    @dtypes(torch.float)
    @parametrize("pass_min_max", [False, True])
    def test_nested_tensor_from_jagged(self, device, dtype, pass_min_max):
        # === construct from (values, offsets) ===
        values = torch.randn(10, 5, device=device, dtype=dtype)
        offsets = torch.tensor([0, 2, 4, 6, 10], device=device, dtype=torch.int64)

        # compute min / max seqlen
        lengths = offsets.diff()
        min_seqlen = lengths.min().item()
        max_seqlen = lengths.max().item()

        if pass_min_max:
            nt = torch.nested.nested_tensor_from_jagged(
                values, offsets=offsets, min_seqlen=min_seqlen, max_seqlen=max_seqlen
            )
        else:
            nt = torch.nested.nested_tensor_from_jagged(values, offsets=offsets)
        self._validate_nt(
            nt,
            device,
            dtype,
            torch.jagged,
            requires_grad=False,
            dim=3,
            batch_size=4,
            contiguous=True,
            cached_min_seqlen=(min_seqlen if pass_min_max else None),
            cached_max_seqlen=(max_seqlen if pass_min_max else None),
            base=values,
        )

        # === construct from (values, offsets, lengths) ===
        lengths = torch.tensor([2, 1, 1, 2], device=device)

        # compute min / max seqlen
        min_seqlen = lengths.min().item()
        max_seqlen = lengths.max().item()

        if pass_min_max:
            nt = torch.nested.nested_tensor_from_jagged(
                values,
                offsets=offsets,
                lengths=lengths,
                min_seqlen=min_seqlen,
                max_seqlen=max_seqlen,
            )
        else:
            nt = torch.nested.nested_tensor_from_jagged(
                values, offsets=offsets, lengths=lengths
            )

        # when both offsets / lengths are specified, expect non-contiguous
        self._validate_nt(
            nt,
            device,
            dtype,
            torch.jagged,
            requires_grad=False,
            dim=3,
            batch_size=4,
            contiguous=False,
            cached_min_seqlen=(min_seqlen if pass_min_max else None),
            cached_max_seqlen=(max_seqlen if pass_min_max else None),
            base=values,
        )
        self.assertIs(nt.lengths(), lengths)

        # === construct from (values, lengths) ===
        values = torch.randn(14, 5, device=device, dtype=dtype)
        lengths = torch.tensor([2, 3, 4, 5], device=device)

        # compute min / max seqlen
        min_seqlen = lengths.min().item()
        max_seqlen = lengths.max().item()

        if pass_min_max:
            nt = torch.nested.nested_tensor_from_jagged(
                values, lengths=lengths, min_seqlen=min_seqlen, max_seqlen=max_seqlen
            )
        else:
            nt = torch.nested.nested_tensor_from_jagged(values, lengths=lengths)

        # for now, if only lengths is specified, convert to offsets to integrate best with the
        # existing kernels
        expected_offsets = torch.tensor([0, 2, 5, 9, 14], device=device)
        expected_nt = torch.nested.nested_tensor_from_jagged(
            values, offsets=expected_offsets
        )
        self._validate_nt(
            nt,
            device,
            dtype,
            torch.jagged,
            requires_grad=False,
            dim=3,
            batch_size=4,
            contiguous=True,
            cached_min_seqlen=(min_seqlen if pass_min_max else None),
            cached_max_seqlen=(max_seqlen if pass_min_max else None),
            base=values,
            ref_nt=expected_nt,
        )

        # error case: no offsets or lengths
        with self.assertRaisesRegex(
            RuntimeError, "At least one of offsets or lengths is required"
        ):
            torch.nested.nested_tensor_from_jagged(values, offsets=None, lengths=None)

    @onlyCPU
    def test_nested_tensor_from_jagged_fx_trace(self, device):
        def fn(x, y):
            return torch.nested.nested_tensor_from_jagged(x, y)

        def user_unwrapped(x, y):
            return fn(x, y)

        with self.assertRaisesRegex(
            RuntimeError,
            "torch.nested.nested_tensor_from_jagged does not support tracing with fx.symbolic_trace",
        ):
            torch.fx.symbolic_trace(user_unwrapped)

    @dtypes(torch.float, torch.double, torch.half)
    @parametrize("dim", range(5))
    @parametrize(
        "layout",
        [torch.strided, torch.jagged],
        name_fn=lambda l: f"layout_{str(l).split('.')[1]}",
    )
    @parametrize("requires_grad", [False, True])
    @parametrize("contiguous", [False, True])
    def test_as_nested_tensor_from_tensor(
        self, device, dtype, dim, layout, requires_grad, contiguous
    ):
        if dim == 0:
            t = torch.tensor(3.0, requires_grad=requires_grad)
        else:
            t = torch.randn(*(3 for _ in range(dim)), requires_grad=requires_grad)
        assert t.dim() == dim

        if dim < 2:
            # 0-1 dim tensors can't be converted to NTs
            with self.assertRaisesRegex(
                RuntimeError, "Expected tensor argument to have dim"
            ):
                nt = torch.nested.as_nested_tensor(
                    t, device=device, dtype=dtype, layout=layout
                )
            return

        orig_t = t
        if not contiguous:
            t = t.transpose(0, 1)

        nt = torch.nested.as_nested_tensor(t, device=device, dtype=dtype, layout=layout)
        expected_dim = t.dim()
        expected_batch_size = t.size(0)
        expected_seqlen = t.size(1) if layout == torch.jagged else None
        self._validate_nt(
            nt,
            device,
            dtype,
            layout,
            requires_grad=requires_grad,
            dim=dim,
            batch_size=expected_batch_size,
            contiguous=True,
            cached_min_seqlen=expected_seqlen,
            cached_max_seqlen=expected_seqlen,
        )

        if torch.device(device) == t.device and dtype == t.dtype and contiguous:
            # should be the non-copying (view) case
            self.assertTrue(nt._is_view() and nt._base is t)

        # should have equivalent components to construction from unbound tensor list
        nt_from_unbind = torch.nested.as_nested_tensor(
            list(t.unbind(0)), device=device, dtype=dtype, layout=layout
        )
        self.assertEqualIgnoringNestedInts(nt, nt_from_unbind)

        # ensure call on a NT with the same properties returns the NT directly
        nt2 = torch.nested.as_nested_tensor(
            nt, device=device, dtype=dtype, layout=layout
        )
        self.assertTrue(nt is nt2)

        # ensure call with device=None uses input tensor device
        nt3 = torch.nested.as_nested_tensor(
            t.to(device=device, dtype=dtype),
            device=None,
            dtype=None,
            layout=layout,
        )
        self._validate_nt(
            nt3,
            device,
            dtype,
            layout,
            requires_grad=requires_grad,
            dim=dim,
            batch_size=expected_batch_size,
            contiguous=True,
            cached_min_seqlen=expected_seqlen,
            cached_max_seqlen=expected_seqlen,
        )

        # we don't support conversion between layouts this way atm
        other_layout = torch.strided if layout == torch.jagged else torch.jagged
        with self.assertRaisesRegex(
            RuntimeError, "Converting between nested tensor layouts is not supported"
        ):
            torch.nested.as_nested_tensor(
                nt, device=device, dtype=dtype, layout=other_layout
            )

        if requires_grad:
            # make sure gradients flow back into inputs
            (nt * 2).backward(torch.ones_like(nt))
            self.assertEqual(orig_t.grad, torch.ones_like(orig_t) * 2)

    @dtypes(torch.double, torch.half)
    @onlyCUDA
    def test_device_dtype_transfer_updates_offsets(self, device, dtype):
        for tensor_list in self._get_example_tensor_lists():
            orig_device = torch.device("cpu")
            orig_dtype = torch.float32
            nt = torch.nested.nested_tensor(
                tensor_list, layout=torch.jagged, device=orig_device, dtype=orig_dtype
            )

            self.assertEqual(torch.int64, nt.offsets().dtype)
            nt = nt.to(device=device).to(dtype=dtype)

            # offsets should still be int64 on the new device
            self.assertEqual(nt.values().device, nt.offsets().device)
            self.assertEqual(torch.int64, nt.offsets().dtype)

    def test_unbind(self, device):
        for tensor_list in self._get_example_tensor_lists():
            nt = torch.nested.nested_tensor(
                tensor_list, layout=torch.jagged, device=device
            )  # ragged_idx = 1
            out = nt.unbind()
            self.assertEqual(len(out), len(tensor_list))
            for i, t in enumerate(out):
                self.assertEqual(t, tensor_list[i])

    @parametrize("ragged_idx", [2, 3])
    def test_unbind_transpose(self, device, ragged_idx):
        for tensor_list in self._get_example_tensor_lists():
            nt = torch.nested.nested_tensor(
                tensor_list, layout=torch.jagged, device=device
            )
            if ragged_idx < nt.dim():
                nt = nt.transpose(1, ragged_idx)  # set ragged_idx
                out = nt.unbind()
                self.assertEqual(len(out), len(tensor_list))
                for i, t in enumerate(out):
                    self.assertEqual(
                        t.transpose(0, ragged_idx - 1), tensor_list[i]
                    )  # transpose back each element of result

    def test_unbind_transpose_ragged_idx_last_dim(self, device):
        for tensor_list in self._get_example_tensor_lists():
            nt = torch.nested.nested_tensor(
                tensor_list, layout=torch.jagged, device=device
            ).transpose(1, -1)  # set ragged_idx = last dimension
            out = nt.unbind()
            self.assertEqual(len(out), len(tensor_list))
            for i, t in enumerate(out):
                self.assertEqual(
                    t.transpose(0, -1), tensor_list[i]
                )  # transpose back each element of result

    def test_unbind_lengths(self, device):
        values = torch.randn(16, 128, device=device)
        offsets = torch.tensor([0, 8, 12, 13, 16], device=device)
        lengths = torch.tensor([6, 2, 1, 2], device=device)
        nt = torch.nested.nested_tensor_from_jagged(
            values, offsets=offsets, lengths=lengths
        )  # 3D nested tensor

        tensor_list = []
        for i in range(offsets.shape[0] - 1):
            tensor_list.append(values[offsets[i] : (offsets[i] + lengths[i])])

        out = nt.unbind()
        self.assertEqual(len(out), len(tensor_list))
        for i, t in enumerate(out):
            self.assertEqual(t, tensor_list[i])

    def test_unbind_lengths_ragged_idx_1(self, device):
        values = torch.randn(16, 8, 128, device=device)
        offsets = torch.tensor([0, 8, 12, 13, 16], device=device)
        lengths = torch.tensor([6, 2, 1, 2], device=device)
        ragged_idx = 1
        nt = torch.nested._internal.nested_tensor.NestedTensor(
            values, offsets=offsets, lengths=lengths, _ragged_idx=ragged_idx
        )  # 4D nested tensor

        tensor_list = []
        for i in range(offsets.shape[0] - 1):
            tensor_list.append(values[offsets[i] : (offsets[i] + lengths[i]), :, :])

        out = nt.unbind()

        self.assertEqual(len(out), len(tensor_list))
        for i, t in enumerate(out):
            self.assertEqual(t, tensor_list[i])

    def test_unbind_lengths_ragged_idx_equals_2_bad_dim(self, device):
        values = torch.randn(16, 8, 128, device=device)
        offsets = torch.tensor([0, 8, 12, 13, 16], device=device)
        lengths = torch.tensor([6, 2, 1, 2], device=device)
        ragged_idx = 2
        nt = torch.nested._internal.nested_tensor.NestedTensor(
            values, offsets=offsets, lengths=lengths, _ragged_idx=ragged_idx
        )  # 4D nested tensor

        self.assertRaisesRegex(
            RuntimeError,
            r"unbind\(\): nested tensor offsets and lengths.*",
            lambda: nt.unbind(),
        )

    def test_unbind_lengths_ragged_idx_2(self, device):
        values = torch.randn(16, 8, 128, device=device)
        offsets = torch.tensor([0, 2, 4, 8], device=device)
        lengths = torch.tensor([2, 1, 3], device=device)
        ragged_idx = 2
        nt = torch.nested._internal.nested_tensor.NestedTensor(
            values, offsets=offsets, lengths=lengths, _ragged_idx=ragged_idx
        )  # 4D nested tensor

        tensor_list = []
        for i in range(offsets.shape[0] - 1):
            tensor_list.append(values[:, offsets[i] : (offsets[i] + lengths[i]), :])

        out = nt.unbind()

        self.assertEqual(len(out), len(tensor_list))
        for i, t in enumerate(out):
            self.assertEqual(t, tensor_list[i])

    def test_unbind_lengths_ragged_idx_3(self, device):
        values = torch.randn(16, 8, 128, device=device)
        offsets = torch.tensor([0, 100, 128], device=device)
        lengths = torch.tensor([50, 28], device=device)
        ragged_idx = 3
        nt = torch.nested._internal.nested_tensor.NestedTensor(
            values, offsets=offsets, lengths=lengths, _ragged_idx=ragged_idx
        )  # 4D nested tensor

        tensor_list = []
        for i in range(offsets.shape[0] - 1):
            tensor_list.append(values[:, :, offsets[i] : (offsets[i] + lengths[i])])

        out = nt.unbind()

        self.assertEqual(len(out), len(tensor_list))
        for i, t in enumerate(out):
            self.assertEqual(t, tensor_list[i])

    @skipIfTorchDynamo(
        "TorchDynamo raises an error for ragged_idx == 0 earlier than Torch"
    )
    def test_unbind_lengths_ragged_idx_0(self, device):
        values = torch.randn(16, 8, 128, device=device)
        offsets = torch.tensor([0, 100, 128], device=device)
        lengths = torch.tensor([50, 28], device=device)
        ragged_idx = 0
        nt = torch.nested._internal.nested_tensor.NestedTensor(
            values, offsets=offsets, lengths=lengths, _ragged_idx=ragged_idx
        )  # 4D nested tensor

        tensor_list = []
        for i in range(offsets.shape[0] - 1):
            tensor_list.append(values[:, :, offsets[i] : (offsets[i] + lengths[i])])

        self.assertRaisesRegex(
            RuntimeError,
            r"unbind\(\): nested tensor.*out of bounds",
            lambda: nt.unbind(),
        )

    def test_narrow(self, device):
        starts = torch.tensor([0, 1, 2, 3, 4], device=device, dtype=torch.int64)
        lengths = torch.tensor([3, 2, 2, 1, 5], device=device, dtype=torch.int64)
        buffer = (
            torch.arange(0, 10, device=device, dtype=torch.int64)
            .unsqueeze(0)
            .expand(5, -1)
            .clone()
            .detach()
        )
        nt = torch.nested.narrow(buffer, 1, starts, lengths, layout=torch.jagged)

        self.assertTrue(nt._is_view() and nt._base is buffer)

        # TODO: Use this approach when unbind is functional
        # unbinded_nt = nt.unbind()
        # for i in range(starts.shape[0]):
        #     self.assertEqual(torch.arange(starts[i], starts[i] + lengths[i], device=device, dtype=torch.int64), unbinded_nt[i])
        for i in range(starts.shape[0]):
            self.assertEqual(
                torch.arange(
                    starts[i], starts[i] + lengths[i], device=device, dtype=torch.int64
                ),
                nt.values()[nt.offsets()[i] : (nt.offsets()[i] + nt.lengths()[i])],
            )

    def test_njt_cat(self, device):
        offsets = torch.tensor([0, 2, 3], device=device, dtype=torch.int64)
        values_1 = torch.randn(
            3, 2, dtype=torch.float64, device=device, requires_grad=True
        )
        values_2 = torch.randn(
            3, 4, dtype=torch.float64, device=device, requires_grad=True
        )

        def grad_test_func(values_1, values_2, offsets):
            nt_1 = torch.nested.nested_tensor_from_jagged(values_1, offsets)
            nt_2 = torch.nested.nested_tensor_from_jagged(values_2, offsets)
            nt_3 = torch.cat([nt_1, nt_2], dim=-1)
            return nt_3.values()

        assert gradcheck(
            grad_test_func,
            inputs=(values_1, values_2, offsets),
            check_batched_grad=False,
        )

    def test_is_contiguous(self, device):
        a = torch.randn(2, 3, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(3, 3, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(4, 3, requires_grad=True, dtype=torch.float64, device=device)
        nt_contiguous = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)

        starts_nc = torch.tensor([0, 1, 2, 3, 4], device=device, dtype=torch.int64)
        lengths_nc = torch.tensor([3, 2, 2, 1, 5], device=device, dtype=torch.int64)
        narrow_base = (
            torch.arange(0, 10, device=device, dtype=torch.int64)
            .unsqueeze(0)
            .expand(5, -1)
            .clone()
        )
        nt_noncontiguous = torch.nested.narrow(
            narrow_base, 1, starts_nc, lengths_nc, layout=torch.jagged
        )

        starts_c = torch.tensor([1, 0, 0, 0, 0], device=device, dtype=torch.int64)
        lengths_c = torch.tensor([9, 10, 10, 10, 8], device=device, dtype=torch.int64)
        nt_contiguous_narrow = torch.nested.narrow(
            narrow_base, 1, starts_c, lengths_c, layout=torch.jagged
        )

        # Test contiguous case
        assert nt_contiguous.is_contiguous()

        # Test narrow case
        assert not nt_noncontiguous.is_contiguous()
        assert nt_contiguous_narrow.is_contiguous()

        # Test querying by memory_format
        self.assertTrue(
            nt_contiguous.is_contiguous(memory_format=torch.contiguous_format)
        )
        self.assertTrue(
            not nt_noncontiguous.is_contiguous(memory_format=torch.contiguous_format)
        )
        self.assertTrue(
            nt_contiguous_narrow.is_contiguous(memory_format=torch.contiguous_format)
        )

    def test_layout_under_torch_dispatch_mode(self):
        from torch.testing._internal.logging_tensor import (
            capture_logs_with_logging_tensor_mode,
        )

        nt = random_nt_from_dims(
            [2, None, 3], torch.device("cpu"), torch.float32, layout=torch.jagged
        )

        with capture_logs_with_logging_tensor_mode():
            self.assertEqual(nt.layout, torch.jagged)

    @skipIfTorchDynamo("Not a suitable test for TorchDynamo")
    @parametrize(
        "func", [torch.empty_like, torch.randn_like], name_fn=lambda f: f.__name__
    )
    def test_like_shape(self, func):
        nt = random_nt_from_dims(
            [2, None, 3], torch.device("cpu"), torch.float32, layout=torch.jagged
        )
        nt_like = func(nt)

        for nt_ub in nt_like.unbind():
            t_like = func(nt_ub)
            self.assertEqual(nt_ub.shape, t_like.shape)

    @skipIfTorchDynamo("Not a suitable test for TorchDynamo")
    @parametrize(
        "func", [torch.ones_like, torch.zeros_like], name_fn=lambda f: f.__name__
    )
    def test_like_value(self, func):
        nt = random_nt_from_dims(
            [2, None, 3], torch.device("cpu"), torch.float32, layout=torch.jagged
        )
        nt_like = func(nt)

        for nt_ub in nt_like.unbind():
            t_like = func(nt_ub)
            self.assertEqual(nt_ub, t_like)

    def test_noncontiguous_pointwise(self, device):
        a = torch.randn(2, 3, 4, requires_grad=True, dtype=torch.float64, device=device)
        b = torch.randn(3, 3, 4, requires_grad=True, dtype=torch.float64, device=device)
        c = torch.randn(4, 3, 4, requires_grad=True, dtype=torch.float64, device=device)
        nt = torch.nested.nested_tensor([a, b, c], layout=torch.jagged)
        # transpose ragged dim
        transposed = nt.transpose(1, 2)
        self.assertFalse(transposed.is_contiguous())
        clone = transposed.clone()

        def check_nt_equality(x, y):
            self.assertEqual(x.values(), y.values())
            self.assertEqual(x.offsets(), y.offsets())
            self.assertEqual(x._ragged_idx, y._ragged_idx)
            self.assertEqual(x.shape, y.shape)

        self.assertFalse(clone.is_contiguous())
        check_nt_equality(clone, transposed)

        clone_contig = transposed.clone(memory_format=torch.contiguous_format)
        self.assertTrue(clone_contig.is_contiguous())
        check_nt_equality(clone_contig, transposed)

        detached = transposed.detach()
        self.assertFalse(clone.is_contiguous())
        check_nt_equality(detached, transposed)

    def test_permute(self, device):
        nt = random_nt_from_dims(
            [2, None, 3, 5], device, torch.float32, layout=torch.jagged
        )
        nt_shape = nt.shape
        nt_inner_shape = nt.values().shape
        with self.assertRaisesRegex(
            ValueError,
            r"permute\(\): number of dimensions in the tensor input \(4\) "
            + r"does not match the length of the desired ordering of dimensions \(3\).",
        ):
            nt.permute(0, 2, 1)
        with self.assertRaisesRegex(
            ValueError, r"permute\(\): duplicate dims are not allowed."
        ):
            nt.permute(0, 2, -2, 3)
        with self.assertRaisesRegex(
            ValueError, "Permute is not supported on the batch dimension for jagged NT"
        ):
            nt.permute(1, 0, 2, 3)
        nt_permute = nt.permute(0, 2, 1, -1)
        self.assertEqual(
            nt_permute.shape, (nt_shape[0], nt_shape[2], nt_shape[1], nt_shape[3])
        )
        self.assertEqual(
            nt_permute.values().shape,
            (nt_inner_shape[1], nt_inner_shape[0], nt_inner_shape[2]),
        )
        self.assertEqual(nt_permute._ragged_idx, 2)
        self.assertEqual(nt_permute.permute(0, 2, 1, 3), nt)

    def test_to_dtype(self, device):
        nt = random_nt_from_dims(
            [2, None, 3], device, torch.float32, layout=torch.jagged
        )
        nt_after = nt.to(torch.float64)
        self.assertEqual(torch.float32, nt.dtype)
        self.assertEqual(torch.float64, nt_after.dtype)
        self.assertEqual(torch.float64, nt_after.values().dtype)
        self.assertEqual(torch.int64, nt_after.offsets().dtype)

        noncontiguous_nt = nt.transpose(1, 2)
        noncontiguous_nt_after = noncontiguous_nt.to(torch.bfloat16)
        self.assertEqual(torch.bfloat16, noncontiguous_nt_after.dtype)
        self.assertEqual(torch.bfloat16, noncontiguous_nt_after.values().dtype)
        self.assertEqual(torch.int64, noncontiguous_nt_after.offsets().dtype)

    def test_to_copy(self, device):
        nt = torch.nested.nested_tensor(
            [
                torch.randn(
                    i + 2, 3, 4, requires_grad=True, dtype=torch.float64, device=device
                )
                for i in range(3)
            ],
            layout=torch.jagged,
        )

        nt_copy_dtype = torch.ops.aten._to_copy(nt, dtype=torch.float16)
        self.assertEqual(torch.float16, nt_copy_dtype.dtype)

        nt_t = nt.transpose(1, 2)
        nt_t_copy_dtype = torch.ops.aten._to_copy(nt_t, dtype=torch.float16)
        self.assertEqual(torch.float16, nt_t_copy_dtype.dtype)

    def test_copy_(self, device):
        offsets = torch.tensor([0, 2, 4], device=device)
        a = torch.nested.nested_tensor_from_jagged(
            torch.zeros(4, 3, device=device), offsets
        )
        b = torch.nested.nested_tensor_from_jagged(
            torch.ones(4, 3, device=device), offsets
        )
        a.copy_(b)
        torch._dynamo.disable(self.assertEqual)(a, b)

        offsets_2 = torch.tensor([0, 2, 4], device=device)
        c = torch.nested.nested_tensor_from_jagged(
            torch.ones(4, 3, device=device), offsets_2
        )
        # fail when tensors have the same size but not the exact same offset tensor.
        with self.assertRaisesRegex(
            RuntimeError,
            "copy_ only supports Nested Tensors that have same size and the exact same offset tensor.",
        ):
            a.copy_(c)

        # fail when tensors have different sizes
        a = a.transpose(1, 2)
        with self.assertRaisesRegex(
            RuntimeError,
            "copy_ only supports Nested Tensors that have same size and the exact same offset tensor.",
        ):
            a.copy_(b)

    @skipIfTorchDynamo("Dynamo doesn't know how to trace prof.events()")
    def test_profiler_sequence_nr(self):
        with torch.profiler.profile() as prof:
            values = torch.randn(4, 6, requires_grad=True)
            offsets = torch.tensor([0, 2, 4])
            values = values * 2
            l = torch.nn.Linear(6, 8)
            nt = torch.nested.nested_tensor_from_jagged(values, offsets)

            nt = l(nt)
            val = nt.values()

            loss = val.sum()
            loss.backward()

        fwd_seq_nrs = []
        for evt in prof.events():
            if (
                "linear" in evt.name.lower()
                and "backward" not in evt.name.lower()
                and evt.sequence_nr != -1
            ):
                fwd_seq_nrs.append(evt.sequence_nr)

        bwd_seq_nrs = []
        for evt in prof.events():
            if (
                "linear" in evt.name.lower()
                and "backward" in evt.name.lower()
                and "evaluate_function" not in evt.name.lower()
                and evt.sequence_nr != -1
            ):
                bwd_seq_nrs.append(evt.sequence_nr)

        # There should only be one such event with a sequence number:
        # the PythonTLSSnapshot event - but, note that it's not terrible if
        # we end up with multiple events with the same sequence number - so we
        # could relax this check if it becomes inconvenient to maintain this
        # property.
        self.assertEqual(len(fwd_seq_nrs), 1)
        self.assertEqual(len(bwd_seq_nrs), 1)
        self.assertEqual(fwd_seq_nrs[0], bwd_seq_nrs[0])

    def test_is_same_size(self, device):
        def get_3_tensors():
            return [
                torch.randn(
                    i + 2, 3, 4, requires_grad=True, dtype=torch.float64, device=device
                )
                for i in range(3)
            ]

        nt1, offsets1 = jagged_from_list(get_3_tensors(), None)
        nt2, offsets1 = jagged_from_list(get_3_tensors(), offsets1)

        nt3, offsets2 = jagged_from_list(get_3_tensors(), None)
        nt4, offsets2 = jagged_from_list(get_3_tensors(), offsets2)

        def check_size(nt1, nt2, nt3, nt4):
            self.assertTrue(torch.ops.aten.is_same_size(nt1, nt2))
            self.assertTrue(torch.ops.aten.is_same_size(nt3, nt4))
            self.assertFalse(torch.ops.aten.is_same_size(nt1, nt3))

        check_size(nt1, nt2, nt3, nt4)

        nt1_t, nt2_t, nt3_t, nt4_t = (x.transpose(1, 2) for x in (nt1, nt2, nt3, nt4))
        check_size(nt1_t, nt2_t, nt3_t, nt4_t)

    @skipIfTorchDynamo("compiles internally")
    @unittest.skipIf(IS_WINDOWS, reason="Windows not yet supported for torch.compile")
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    def test_specialize_dynamic_shape(self, device):
        values = torch.randn((18, 16), device=device)
        offsets = torch.tensor([0, 2, 3, 6, 15, 18], device=device)
        like_values = torch.randn_like(values)

        # this marks values as dynamic
        nt = torch.nested.nested_tensor_from_jagged(values, offsets)

        def fn(values, same_size):
            # here, the dynamic shape is specialized by same_size's shape
            # https://github.com/pytorch/pytorch/issues/127097
            # make sure this doesn't error out in torch.compile
            return values + same_size

        self.assertEqual(
            fn(values, like_values),
            torch.compile(fn)(values, like_values),
        )

    @skipIfTorchDynamo("compiles internally")
    @unittest.skipIf(IS_WINDOWS, reason="Windows not yet supported for torch.compile")
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    def test_specialize_dynamic_shape_recompile(self, device):
        def generate_inp(total_len):
            values = torch.randn((total_len, 16), device=device)
            offsets = torch.tensor([0, 2, 3, 6, 15, total_len], device=device)
            like_values = torch.randn_like(values)
            return values, offsets, like_values

        def check_results(ref_fn, res_fn, args):
            values, offsets, like_values = args
            # this may add dynamic shape markings
            # goal of this test is to make sure that whatever markings are there,
            # we eventually stop recompiling as shape changes.
            nt = torch.nested.nested_tensor_from_jagged(values, offsets)

            self.assertEqual(ref_fn(values, like_values), res_fn(values, like_values))

        def fn(values, same_size):
            return values + same_size

        compile_counter = torch._dynamo.testing.CompileCounter()

        compiled_fn = torch._dynamo.optimize(compile_counter, nopython=True)(fn)
        check_results(fn, compiled_fn, generate_inp(18))
        self.assertEqual(compile_counter.frame_count, 1)

        check_results(fn, compiled_fn, generate_inp(19))
        # we'll probably recompile here with dynamic shapes - it's okay if not though.
        frame_count_2 = compile_counter.frame_count
        self.assertIn(frame_count_2, [1, 2])

        # make sure that by now we've already compiled with dynamic shapes, so additional
        # shapes should not trigger additional recompiles.
        check_results(fn, compiled_fn, generate_inp(20))
        self.assertEqual(compile_counter.frame_count, frame_count_2)

    # Note 1: Math fallback doesn't work with bfloat16 on CUDA
    # Note 2: ROCm doesn't support flash attention or mem_efficient attention for NT
    @unittest.skipIf(
        TEST_WITH_ROCM,
        "ROCm doesn't support flash attention or mem_efficient attention for NT",
    )
    @dtypes(
        *(
            [torch.float16, torch.bfloat16, torch.float32]
            if SM80OrLater
            else [torch.float16, torch.float32]
        )
    )
    def test_sdpa(self, device, dtype):
        batch_size = 1
        emb_dims = 128
        n_heads = 8
        head_dims = emb_dims // n_heads

        sen1 = torch.randn(11, emb_dims, dtype=dtype, device=device)
        sen2 = torch.randn(13, emb_dims, dtype=dtype, device=device)

        query = torch.nn.Linear(
            emb_dims, emb_dims, bias=False, device=device, dtype=dtype
        )
        key = torch.nn.Linear(
            emb_dims, emb_dims, bias=False, device=device, dtype=dtype
        )
        value = torch.nn.Linear(
            emb_dims, emb_dims, bias=False, device=device, dtype=dtype
        )

        # Simplest case: 1 sentence, no batching
        x_d1 = sen1.unsqueeze(0)
        x_nt = torch.nested.as_nested_tensor([sen1], layout=torch.jagged)

        # See note below for why we detach here.
        q_d1 = (
            query(x_d1)
            .view(batch_size, -1, n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        q_d1_t = q_d1.transpose(1, 2)
        k_d1 = (
            key(x_d1)
            .view(batch_size, -1, n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        k_d1_t = k_d1.transpose(1, 2)
        v_d1 = (
            value(x_d1)
            .view(batch_size, -1, n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        v_d1_t = v_d1.transpose(1, 2)

        q_nt = (
            query(x_nt)
            .view(*x_nt.size()[0:2], n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        q_nt_t = q_nt.transpose(1, 2)
        k_nt = (
            key(x_nt)
            .view(*x_nt.size()[0:2], n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        k_nt_t = k_nt.transpose(1, 2)
        v_nt = (
            value(x_nt)
            .view(*x_nt.size()[0:2], n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        v_nt_t = v_nt.transpose(1, 2)

        # High Precision Math Reference
        q_d1_f32 = q_d1.to(torch.float32)
        k_d1_f32 = k_d1.to(torch.float32)
        v_d1_f32 = v_d1.to(torch.float32)
        q_d1_f32_t = q_d1_f32.transpose(1, 2)
        k_d1_f32_t = k_d1_f32.transpose(1, 2)
        v_d1_f32_t = v_d1_f32.transpose(1, 2)
        out_ref = torch.ops.aten._scaled_dot_product_attention_math(
            q_d1_f32_t, k_d1_f32_t, v_d1_f32_t
        )[0]
        grads_ref = torch.autograd.grad(out_ref.sum(), (q_d1_f32, k_d1_f32, v_d1_f32))

        # Low Precision Math Reference
        out_lp_ref = torch.ops.aten._scaled_dot_product_attention_math(
            q_d1_t, k_d1_t, v_d1_t
        )[0]
        grads_lp_ref = torch.autograd.grad(out_lp_ref.sum(), (q_d1, k_d1, v_d1))

        # Compute tolerances
        output_ref_atol, output_ref_rtol = get_tolerances(out_ref, out_lp_ref)
        grad_q_ref_atol, grad_q_ref_rtol = get_tolerances(grads_ref[0], grads_lp_ref[0])
        grad_k_ref_atol, grad_k_ref_rtol = get_tolerances(grads_ref[1], grads_lp_ref[1])
        grad_v_ref_atol, grad_v_ref_rtol = get_tolerances(grads_ref[2], grads_lp_ref[2])
        grad_atols = [grad_q_ref_atol, grad_k_ref_atol, grad_v_ref_atol]
        grad_rtols = [grad_q_ref_rtol, grad_k_ref_rtol, grad_v_ref_rtol]

        attn_d1 = torch.nn.functional.scaled_dot_product_attention(
            q_d1_t, k_d1_t, v_d1_t
        ).transpose(1, 2)
        attn_nt = torch.nn.functional.scaled_dot_product_attention(
            q_nt_t, k_nt_t, v_nt_t
        ).transpose(1, 2)

        self.assertEqual(
            attn_d1,
            attn_nt.unbind()[0].unsqueeze(0),
            atol=output_ref_atol,
            rtol=output_ref_rtol,
        )

        # Simple case: 2 sentences, no extra params
        x_d2 = sen2.unsqueeze(0)
        x_nt = torch.nested.as_nested_tensor([sen1, sen2], layout=torch.jagged)

        # NB: we make sure the leaf tensor we compute gradients for is the view-ed tensor before
        # it is transposed. This is because today we cannot backward through view or unbind a
        # transposed tensor.
        q_d2 = (
            query(x_d2)
            .view(batch_size, -1, n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        q_d2_t = q_d2.transpose(1, 2)
        k_d2 = (
            key(x_d2)
            .view(batch_size, -1, n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        k_d2_t = k_d2.transpose(1, 2)
        v_d2 = (
            value(x_d2)
            .view(batch_size, -1, n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        v_d2_t = v_d2.transpose(1, 2)

        q_nt = (
            query(x_nt)
            .view(*x_nt.size()[0:2], n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        q_nt_t = q_nt.transpose(1, 2)
        k_nt = (
            key(x_nt)
            .view(*x_nt.size()[0:2], n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        k_nt_t = k_nt.transpose(1, 2)
        v_nt = (
            value(x_nt)
            .view(*x_nt.size()[0:2], n_heads, head_dims)
            .detach()
            .requires_grad_(True)
        )
        v_nt_t = v_nt.transpose(1, 2)

        attn_d2 = torch.nn.functional.scaled_dot_product_attention(
            q_d2_t, k_d2_t, v_d2_t
        ).transpose(1, 2)
        d1_grads = torch.autograd.grad(attn_d1.sum(), (q_d1, k_d1, v_d1))
        d2_grads = torch.autograd.grad(attn_d2.sum(), (q_d2, k_d2, v_d2))

        # Simple case 3: batch_size = 1, seq_len = 1
        q_3 = torch.randn(1, 8, 16, dtype=dtype, device=device)
        q_nt_3 = torch.nested.as_nested_tensor([q_3], layout=torch.jagged)
        q_nt_3 = q_nt_3.transpose(1, 2)
        attn_out = torch.nn.functional.scaled_dot_product_attention(
            q_nt_3, q_nt_3, q_nt_3
        )
        self.assertEqual(attn_out.shape, q_nt_3.shape)

        def check_forward_backward():
            attn_nt = torch.nn.functional.scaled_dot_product_attention(
                q_nt_t, k_nt_t, v_nt_t
            ).transpose(1, 2)

            attn_nts = attn_nt.unbind()
            self.assertEqual(
                attn_d1,
                attn_nts[0].unsqueeze(0),
                atol=output_ref_atol,
                rtol=output_ref_rtol,
            )
            self.assertEqual(
                attn_d2,
                attn_nts[1].unsqueeze(0),
                atol=output_ref_atol,
                rtol=output_ref_rtol,
            )

            nt_grads = torch.autograd.grad(attn_nt.values().sum(), (q_nt, k_nt, v_nt))
            for nt_grad, d1_grad, d2_grad, grad_atol, grad_rtol in zip(
                nt_grads, d1_grads, d2_grads, grad_atols, grad_rtols
            ):
                unbound_nt_grads = nt_grad.unbind()
                self.assertEqual(
                    d1_grad,
                    unbound_nt_grads[0].unsqueeze(0),
                    atol=grad_atol,
                    rtol=grad_rtol,
                )
                self.assertEqual(
                    d2_grad,
                    unbound_nt_grads[1].unsqueeze(0),
                    atol=grad_atol,
                    rtol=grad_rtol,
                )

        # Default
        check_forward_backward()

        # Test dispatcher works by calling only mem-effn and math (as they are safe for all devices)
        with torch.backends.cuda.sdp_kernel(
            enable_flash=False, enable_mem_efficient=True, enable_math=True
        ):
            check_forward_backward()

        # Test math fallback
        with torch.backends.cuda.sdp_kernel(
            enable_flash=False, enable_mem_efficient=False, enable_math=True
        ):
            # Math fallback doesn't work with bfloat16 on CUDA because
            # "group_gemm_dispatch" not implemented for 'BFloat16'
            if not (str(device).startswith("cuda") and dtype == torch.bfloat16):
                check_forward_backward()

    @skipIfTorchDynamo("SDPA test compiles internally")
    @unittest.skipIf(IS_WINDOWS, reason="Windows not yet supported for torch.compile")
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    # Guarding with sqrt() doesn't work on ROCm?
    @skipCUDAIfRocm
    @onlyCUDA
    @dtypes(
        *(
            [torch.float16, torch.bfloat16, torch.float32]
            if SM80OrLater
            else [torch.float16, torch.float32]
        )
    )
    def test_sdpa_compile(self, device, dtype):
        batch_size = 1
        emb_dims = 1024
        n_heads = 8
        head_dims = emb_dims // n_heads

        sen1 = torch.randn(11, emb_dims, dtype=dtype, device=device)
        sen2 = torch.randn(13, emb_dims, dtype=dtype, device=device)

        query = torch.nn.Linear(
            emb_dims, emb_dims, bias=False, device=device, dtype=dtype
        )
        key = torch.nn.Linear(
            emb_dims, emb_dims, bias=False, device=device, dtype=dtype
        )
        value = torch.nn.Linear(
            emb_dims, emb_dims, bias=False, device=device, dtype=dtype
        )

        # Simplest case: 1 sentence, no batching
        x_d1 = sen1.unsqueeze(0)
        x_d2 = sen2.unsqueeze(0)
        x_nt = torch.nested.as_nested_tensor([sen1, sen2], layout=torch.jagged)

        q_d1 = query(x_d1).view(batch_size, -1, n_heads, head_dims).transpose(1, 2)
        k_d1 = key(x_d1).view(batch_size, -1, n_heads, head_dims).transpose(1, 2)
        v_d1 = value(x_d1).view(batch_size, -1, n_heads, head_dims).transpose(1, 2)
        q_d2 = query(x_d2).view(batch_size, -1, n_heads, head_dims).transpose(1, 2)
        k_d2 = key(x_d2).view(batch_size, -1, n_heads, head_dims).transpose(1, 2)
        v_d2 = value(x_d2).view(batch_size, -1, n_heads, head_dims).transpose(1, 2)

        q_nt = (
            query(x_nt)
            .view(*x_nt.size()[0:2], n_heads, head_dims)
            .detach()
            .transpose(1, 2)
        )
        k_nt = (
            key(x_nt)
            .view(*x_nt.size()[0:2], n_heads, head_dims)
            .detach()
            .transpose(1, 2)
        )
        v_nt = (
            value(x_nt)
            .view(*x_nt.size()[0:2], n_heads, head_dims)
            .detach()
            .transpose(1, 2)
        )

        # High Precision Math Reference
        q_d1_f32 = q_d1.to(torch.float32)
        k_d1_f32 = k_d1.to(torch.float32)
        v_d1_f32 = v_d1.to(torch.float32)
        out_ref = torch.ops.aten._scaled_dot_product_attention_math(
            q_d1_f32, k_d1_f32, v_d1_f32
        )[0]
        # Low Precision Math Reference
        out_lp_ref = torch.ops.aten._scaled_dot_product_attention_math(
            q_d1, k_d1, v_d1
        )[0]
        output_ref_atol, output_ref_rtol = get_tolerances(out_ref, out_lp_ref)

        attn_d1 = torch.nn.functional.scaled_dot_product_attention(
            q_d1, k_d1, v_d1
        ).transpose(1, 2)
        attn_d2 = torch.nn.functional.scaled_dot_product_attention(
            q_d2, k_d2, v_d2
        ).transpose(1, 2)

        compiled_sdpa = torch.compile(torch.nn.functional.scaled_dot_product_attention)
        attn_nt = compiled_sdpa(q_nt, k_nt, v_nt).transpose(1, 2)

        attn_nts = attn_nt.unbind()
        self.assertEqual(
            attn_d1,
            attn_nts[0].unsqueeze(0),
            atol=output_ref_atol,
            rtol=output_ref_rtol,
        )
        self.assertEqual(
            attn_d2,
            attn_nts[1].unsqueeze(0),
            atol=output_ref_atol,
            rtol=output_ref_rtol,
        )

    @dtypes(torch.float32, torch.double, torch.half)
    def test_sdpa_with_constant_sequence_length(self, device, dtype):
        # shape (B, P*, S, D)
        # B: batch size
        # P*: ragged number of prompts
        # S: (constant) sequence length
        # D: embedding size
        query = random_nt_from_dims(
            [4, None, 8, 10],
            device=device,
            dtype=dtype,
            layout=torch.jagged,
            requires_grad=True,
        )
        key = random_nt_from_similar(query)
        value = random_nt_from_similar(query)
        output = F.scaled_dot_product_attention(query, key, value)
        self.assertTrue(isinstance(output, NestedTensor))
        output.values().sum().backward()

        query_dense = query.clone().detach().requires_grad_(True)
        # should be equivalent to just running the buffers through
        output_dense = F.scaled_dot_product_attention(
            query_dense.values(), key.values(), value.values()
        )
        torch._dynamo.disable(self.assertEqual)(output._values, output_dense)
        output_dense.sum().backward()
        torch._dynamo.disable(self.assertEqual)(query.grad, query_dense.grad)

    @onlyCUDA
    @unittest.skipIf(
        not PLATFORM_SUPPORTS_FUSED_ATTENTION,
        "Platform doesn't support flash or mem-efficient attention",
    )
    @dtypes(
        *(
            [torch.float16, torch.bfloat16, torch.float32]
            if SM80OrLater
            else [torch.float16, torch.float32]
        )
    )
    def test_sdpa_with_packed_in_proj(self, device, dtype):
        # shape (B, *, D)
        input_packed = random_nt_from_dims(
            [5, None, 10], device=device, dtype=dtype, layout=torch.jagged
        )

        # Do input projection.
        num_heads = 2
        # should be multiple of 4 for efficient kernels (e.g. flash / mem-efficient)
        head_dim = 8
        qkv_linear = torch.nn.Linear(10, num_heads * head_dim * 3).to(
            device=device, dtype=dtype
        )

        def in_proj(input_packed, qkv_linear=qkv_linear):
            qkv_post_proj = qkv_linear(input_packed)
            # these are non-contiguous to trigger _is_safe_to_get_storage_as_tensor()
            q, k, v = qkv_post_proj.chunk(3, dim=-1)
            q = q.unflatten(-1, [num_heads, head_dim]).transpose(-2, -3)
            k = k.unflatten(-1, [num_heads, head_dim]).transpose(-2, -3)
            v = v.unflatten(-1, [num_heads, head_dim]).transpose(-2, -3)
            return q, k, v

        q, k, v = in_proj(input_packed)
        output = F.scaled_dot_product_attention(q, k, v, attn_mask=None)

        # compare to individually running unbound components through
        for in_component, out_component in zip(
            input_packed.unbind(), output.transpose(-2, -3).unbind()
        ):
            q, k, v = in_proj(in_component)
            out = F.scaled_dot_product_attention(q, k, v).transpose(-2, -3)

            # Low Precision Math Reference
            out_lp_ref = torch.ops.aten._scaled_dot_product_attention_math(q, k, v)[
                0
            ].transpose(-2, -3)
            output_ref_atol, output_ref_rtol = get_tolerances(
                out, out_lp_ref, fudge_factor=2
            )

            self.assertEqual(
                out, out_component, atol=output_ref_atol, rtol=output_ref_rtol
            )

    @skipIfTorchDynamo("SDPA test compiles internally")
    @unittest.skipIf(IS_WINDOWS, reason="Windows not yet supported for torch.compile")
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    # mha_varlen_fwd not supported on ROCm
    @skipCUDAIfRocm
    @onlyCUDA
    @dtypes(
        *(
            [torch.float16, torch.bfloat16, torch.float32]
            if SM80OrLater
            else [torch.float16, torch.float32]
        )
    )
    def test_sdpa_backwards(self, device, dtype):
        values = torch.randn(9, 3, 256, requires_grad=True, device=device, dtype=dtype)
        offsets = torch.tensor([0, 1, 3, 5, 9], device=device, dtype=torch.int64)

        @torch.compile
        def f(values, offsets):
            nt = convert_jagged_to_nested_tensor(values, offsets, max_length=4)
            nt = nt.transpose(-2, -3)
            # purposefully graph break to trigger view replay for subclass view input
            torch.tensor(1).item()
            output = F.scaled_dot_product_attention(nt, nt, nt).transpose(-2, -3)
            return convert_nt_to_jagged(output)

        output = f(values, offsets)
        output.sum().backward()
        self.assertEqual(values.grad, torch.ones_like(values))

    @unittest.skipIf(
        not PLATFORM_SUPPORTS_FUSED_ATTENTION,
        "Platform doesn't support flash or mem-efficient attention",
    )
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    @skipCUDAIfRocm
    @onlyCUDA
    @skipIfTorchDynamo()
    @unittest.skipIf(IS_WINDOWS, reason="Windows not yet supported for torch.compile")
    def test_sdpa_autocast(self, device):
        def fn_nt(values32, values16, offsets):
            nt32 = convert_jagged_to_nested_tensor(values32, offsets, max_length=16)
            nt16 = convert_jagged_to_nested_tensor(values16, offsets, max_length=16)
            nt32 = nt32.transpose(1, 2)
            nt16 = nt16.transpose(1, 2)
            return F.scaled_dot_product_attention(nt32, nt16, nt32)

        def fn_dense(x32, x16):
            x32 = x32.view(8, 16, 4, 16).transpose(1, 2)
            x16 = x16.view(8, 16, 4, 16).transpose(1, 2)
            return F.scaled_dot_product_attention(x32, x16, x32)

        values32 = torch.randn((8 * 16, 4, 16), device=device, dtype=torch.float32)
        values16 = torch.randn((8 * 16, 4, 16), device=device, dtype=torch.float16)
        offsets = torch.arange(0, 8 * 16 + 1, 16, device=device, dtype=torch.int32)

        x32 = values32.clone()
        x16 = values16.clone()

        with torch.autocast(device_type="cuda", dtype=torch.float16):
            out_dense_eager = fn_dense(x32, x16)
            out_dense_compiled = torch.compile(fn_dense)(x32, x16)
            out_nt_eager = fn_nt(values32, values16, offsets)
            out_nt_compiled = torch.compile(fn_nt)(values32, values16, offsets)

        self.assertEqual(out_dense_eager, out_dense_compiled)
        self.assertEqual(
            out_dense_eager.transpose(1, 2),
            out_nt_eager.values().transpose(0, 1).view(8, 16, 4, 16),
        )
        self.assertEqual(
            out_dense_eager.transpose(1, 2),
            out_nt_compiled.values().transpose(0, 1).view(8, 16, 4, 16),
        )

        def get_values():
            return tuple(
                x.clone().detach().requires_grad_(True) for x in (values32, values16)
            )

        v32_dense_eager, v16_dense_eager = get_values()
        v32_dense_compile, v16_dense_compile = get_values()
        v32_nt_eager, v16_nt_eager = get_values()
        v32_nt_compile, v16_nt_compile = get_values()

        with torch.autocast(device_type="cuda", dtype=torch.float16):
            loss_dense_eager = fn_dense(v32_dense_eager, v16_dense_eager).sum()
            loss_dense_compile = torch.compile(fn_dense)(
                v32_dense_compile, v16_dense_compile
            ).sum()
            loss_nt_eager = fn_nt(v32_nt_eager, v16_nt_eager, offsets).values().sum()
            loss_nt_compile = (
                torch.compile(fn_nt)(v32_nt_compile, v16_nt_compile, offsets)
                .values()
                .sum()
            )

        loss_dense_eager.backward()
        loss_dense_compile.backward()
        loss_nt_eager.backward()
        loss_nt_compile.backward()

        self.assertEqual(v32_dense_eager.grad, v32_dense_compile.grad)
        self.assertEqual(v32_dense_eager.grad, v32_nt_eager.grad)
        self.assertEqual(v32_dense_eager.grad, v32_nt_compile.grad)

        self.assertEqual(v16_dense_eager.grad, v16_dense_compile.grad)
        self.assertEqual(v16_dense_eager.grad, v16_nt_eager.grad)
        self.assertEqual(v16_dense_eager.grad, v16_nt_compile.grad)

    @unittest.skipIf(
        not PLATFORM_SUPPORTS_FUSED_ATTENTION,
        "Platform doesn't support flash or mem-efficient attention",
    )
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    @skipCUDAIfRocm
    @onlyCUDA
    @skipIfTorchDynamo()
    def test_sdpa_flop_counter(self, device):
        from torch.utils.flop_counter import FlopCounterMode

        def get_flops(nt):
            flop_counter = FlopCounterMode(display=False)
            with flop_counter:
                ret = torch.nn.functional.scaled_dot_product_attention(nt, nt, nt)
                ret.values().sum().backward()
            return flop_counter.get_total_flops()

        values = torch.randn(
            (8 * 16, 4, 16), requires_grad=True, device=device, dtype=torch.float16
        )
        offsets = torch.arange(0, 8 * 16 + 1, 16, device=device, dtype=torch.int32)
        nt = convert_jagged_to_nested_tensor(values, offsets, max_length=16)

        values_meta = torch.randn(
            (8 * 16, 4, 16), requires_grad=True, device="meta", dtype=torch.float16
        )
        offsets_meta = torch.arange(0, 8 * 16 + 1, 16, device="meta", dtype=torch.int32)
        nt_meta = convert_jagged_to_nested_tensor(values, offsets, max_length=16)

        self.assertEqual(get_flops(nt), get_flops(nt_meta))

    @skipIfTorchDynamo()
    def test_nested_tensor_activation_checkpoint(self, device):
        values = torch.randn(
            9, 3, 256, requires_grad=True, device=device, dtype=torch.float32
        )
        lengths = torch.tensor([1, 2, 3, 3], device=device, dtype=torch.int64)
        offsets = F.pad(lengths, pad=(1, 0)).cumsum(dim=0)

        def fn(values, offsets):
            nt = convert_jagged_to_nested_tensor(values, offsets, max_length=4)
            return convert_nt_to_jagged(nt).sum()

        checkpoint(fn, values, offsets, use_reentrant=False).backward()
        self.assertIsNotNone(values.grad)

        context_fn = partial(
            create_selective_checkpoint_contexts, [torch.ops.aten.cumsum.default]
        )

        values.grad = None

        def fn(values, lengths):
            offsets = F.pad(lengths, pad=(1, 0)).cumsum(dim=0)
            nt = convert_jagged_to_nested_tensor(values, offsets, max_length=4)
            return convert_nt_to_jagged(nt).sum()

        checkpoint(
            fn, values, lengths, use_reentrant=False, context_fn=context_fn
        ).backward()
        self.assertIsNotNone(values.grad)

    # Internally-defined NT use cases are lifted to here for maximum test realism.
    # TODO: Remove these when ViewNestedFromBuffer, etc. are deprecated.
    @skipCUDAIfRocm  # not needed
    @skipIfTorchDynamo("compiles internally")
    @unittest.skipIf(IS_WINDOWS, reason="Windows not yet supported for torch.compile")
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    @parametrize("use_legacy_api", [True, False])
    @skipCPUIf(True, "SPDA Math NT fallback causes failure: see issue #133644")
    def test_dummy_mha_with_nt(self, device, use_legacy_api):
        bs = 3
        d1 = 2
        d2 = 4
        d3 = 16
        n_heads = 2
        d_head = d3 // n_heads
        max_length_1 = 10
        max_length_2 = 20
        torch.manual_seed(0)

        class mha(torch.nn.Module):
            def __init__(self, use_legacy_api) -> None:
                super().__init__()
                torch.manual_seed(0)
                self.linear = torch.nn.Linear(d2, d3, device=device)
                self.use_legacy_api = use_legacy_api

            def forward(self, query, value, offsets):
                value = self.linear(value)
                if self.use_legacy_api:
                    key = convert_jagged_to_nested_tensor_legacy(
                        value, offsets, max_length_1
                    )
                    value = convert_jagged_to_nested_tensor_legacy(
                        value, offsets, max_length_2
                    )
                    query = convert_dense_to_nested_tensor_legacy(query)
                else:
                    key = convert_jagged_to_nested_tensor(value, offsets, max_length_1)
                    value = convert_jagged_to_nested_tensor(
                        value, offsets, max_length_2
                    )
                    query = convert_dense_to_nested_tensor(query)
                q = query.view(bs, -1, n_heads, d_head).transpose(1, 2)
                k = key.view(bs, -1, n_heads, d_head).transpose(1, 2)
                v = value.view(bs, -1, n_heads, d_head).transpose(1, 2)

                with torch.nn.attention.sdpa_kernel(
                    [
                        torch.nn.attention.SDPBackend.FLASH_ATTENTION,
                        torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION,
                    ]
                ):
                    attn_output = torch.nn.functional.scaled_dot_product_attention(
                        q,
                        k,
                        v,
                        attn_mask=None,
                        dropout_p=0.0,
                        is_causal=False,
                    )
                attn_output = attn_output.transpose(1, 2)
                if self.use_legacy_api:
                    attn_output = convert_nt_to_jagged_legacy(attn_output)
                else:
                    attn_output = convert_nt_to_jagged(attn_output)
                return attn_output, key._max_seqlen, value._max_seqlen

        query = torch.rand(bs, d1, d3, device=device)
        value = torch.rand(30, d2, requires_grad=True, device=device)
        # total_length must > than max_length otherwise flash_attn backwark will fail
        offsets = torch.tensor([0, 2, 3, 30], device=device)

        m = mha(use_legacy_api)
        symbolic_traced: torch.fx.GraphModule = torch.fx.symbolic_trace(m)
        m = torch.compile(symbolic_traced)
        attn_output, cached_key_max_seqlen, cached_value_max_seqlen = m(
            query, value, offsets
        )
        loss = attn_output.sum()
        # Check that NT can be fx traced and torch.compile, and backward works
        loss.backward()

        # Check that value.requires_grad is not lost after tracing and compiling
        value_grad = value.grad  # save for comparison later
        self.assertIsNotNone(value_grad)
        # check that max_seqlen is cached properly
        self.assertEqual(cached_key_max_seqlen, max_length_1)
        self.assertEqual(cached_value_max_seqlen, max_length_2)

        # check if the output is numerically equivalent with the eager mode
        m_eager = mha(use_legacy_api)

        value.grad = None
        attn_output_eager, _, _ = m_eager(query, value, offsets)
        attn_output_eager.sum().backward()
        self.assertTrue(torch.allclose(attn_output_eager, attn_output))
        self.assertTrue(torch.allclose(value_grad, value.grad))

    @dtypes(torch.float32)
    def test_apply_(self, device, dtype):
        nt = random_nt_from_dims(
            [5, None, 10],
            device=device,
            dtype=dtype,
            layout=torch.jagged,
            requires_grad=True,
        )

        def f(x):
            return x * 2

        if device != "cpu":
            with self.assertRaisesRegex(
                TypeError, "apply_ is only implemented on CPU tensors"
            ):
                nt.apply_(f)
            return

        before = nt._values.clone().detach()

        nt.apply_(f)
        expected = f(before)
        self.assertEqual(expected, nt._values)
        # apply_ should swap values in-place without appending to autograd graph
        self.assertIsNone(nt.grad)
        self.assertIsNone(nt._values.grad_fn)

    @dtypes(torch.float64, torch.float32, torch.half)
    def test_jagged_padded_dense_conversion_kernels(self, device, dtype):
        values = torch.randn(10, 5, device=device, dtype=dtype)
        offsets = torch.tensor([0, 1, 3, 8, 10], device=device, dtype=torch.int64)
        max_length = offsets.diff().max().item()
        padding_value = 1.3

        # convert jagged -> padded dense
        padded = torch.ops.aten._jagged_to_padded_dense_forward(
            values, [offsets], [max_length], padding_value
        )

        batch_size = offsets.shape[0] - 1
        expected_padded_shape = (batch_size, max_length, values.shape[-1])
        self.assertEqual(padded.shape, expected_padded_shape)

        # convert padded dense -> jagged
        total_L = values.shape[0]
        output_jagged = torch.ops.aten._padded_dense_to_jagged_forward(
            padded, [offsets], total_L
        )

        # should be equivalent to the original values
        self.assertEqual(values, output_jagged)

        # success case: truncate to max length as needed
        trunc_max_length = max_length - 1
        trunc_padded = torch.ops.aten._jagged_to_padded_dense_forward(
            values, [offsets], [trunc_max_length], padding_value
        )
        self.assertEqual(padded[:, :trunc_max_length, :], trunc_padded)

        # specific to CPU impls
        if device == "cpu":
            # error case: multiple offsets on cpu since CPU kernels don't support more now
            with self.assertRaisesRegex(
                RuntimeError, "only a single jagged dim is supported"
            ):
                torch.ops.aten._jagged_to_padded_dense_forward(
                    values, [offsets, offsets], [max_length, max_length], padding_value
                )

            with self.assertRaisesRegex(
                RuntimeError, "only a single jagged dim is supported"
            ):
                torch.ops.aten._padded_dense_to_jagged_forward(
                    padded, [offsets, offsets], total_L
                )

            # error case: > 1D offsets
            offsets2d = offsets.unsqueeze(-1)
            with self.assertRaisesRegex(RuntimeError, "expected 1D offsets"):
                torch.ops.aten._jagged_to_padded_dense_forward(
                    values, [offsets2d], [max_length], padding_value
                )

            with self.assertRaisesRegex(RuntimeError, "expected 1D offsets"):
                torch.ops.aten._padded_dense_to_jagged_forward(
                    padded, [offsets2d], total_L
                )

            # error case: final offset != total_L
            offsets_wrong = offsets.clone().detach()
            offsets_wrong[-1] = total_L + 1
            with self.assertRaisesRegex(
                RuntimeError, "final offset should match total_L value"
            ):
                torch.ops.aten._padded_dense_to_jagged_forward(
                    padded, [offsets_wrong], total_L
                )

            # error case: 1D padded input
            padded_wrong = padded.flatten().clone().detach()
            with self.assertRaisesRegex(RuntimeError, "expected padded dim >= 2"):
                torch.ops.aten._padded_dense_to_jagged_forward(
                    padded_wrong, [offsets], total_L
                )

            # error case: batch item has length > max length
            # max_length is 5 above; 7 here
            offsets_wrong = torch.tensor(
                [0, 1, 8, 9, 10], device=device, dtype=torch.int64
            )
            with self.assertRaisesRegex(RuntimeError, "found batch item of length"):
                torch.ops.aten._padded_dense_to_jagged_forward(
                    padded, [offsets_wrong], total_L
                )

    @dtypes(torch.float32)
    @skipIfTorchDynamo("Test compiles internally")
    @unittest.skipIf(
        sys.version_info >= (3, 12), "torch.compile is not supported on python 3.12+"
    )
    @unittest.skipIf(IS_WINDOWS, reason="Windows not yet supported for torch.compile")
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    @skipCUDAIfRocm
    def test_compile_preserves_metadata_cache(self, device, dtype):
        # shape (B, *, D)
        nt = random_nt_from_dims(
            [4, None, 3, 16],
            device=device,
            dtype=dtype,
            layout=torch.jagged,
            requires_grad=True,
        )

        # expect min / max seqlen to be stored here
        cache = dict(nt._metadata_cache)

        @torch.compile
        def f(nt):
            q = nt.transpose(-3, -2)
            output = F.scaled_dot_product_attention(q, q, q).transpose(-3, -2)
            return output

        output = f(nt)
        output.backward(torch.ones_like(output))
        self.assertEqual(output._metadata_cache, cache)

    @dtypes(torch.float32)
    @skipIfTorchDynamo("Test compiles internally")
    @unittest.skipIf(
        sys.version_info >= (3, 12), "torch.compile is not supported on python 3.12+"
    )
    @unittest.skipIf(IS_WINDOWS, reason="Windows not yet supported for torch.compile")
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    @skipCUDAIfRocm
    def test_compile_with_dynamic_max_seq_len(self, device, dtype):
        # shape (B, *, D)
        # max seq len: 18
        nt = torch.nested.nested_tensor(
            [
                torch.randn(2, 5),
                torch.randn(3, 5),
                torch.randn(18, 5),
            ],
            layout=torch.jagged,
        )

        # max seq len: 19
        nt2 = torch.nested.nested_tensor(
            [
                torch.randn(2, 5),
                torch.randn(3, 5),
                torch.randn(19, 5),
            ],
            layout=torch.jagged,
        )

        def f(nt):
            # TODO: Replace with public API when we can use @properties
            return torch.ones_like(nt) * nt._get_max_seqlen()

        for dynamic in [False, True, None]:
            self.assertFalse(_recompiles_for_inputs(f, (nt,), (nt2,), dynamic=dynamic))

    @dtypes(torch.float32)
    @skipIfTorchDynamo("Test compiles internally")
    @unittest.skipIf(
        sys.version_info >= (3, 12), "torch.compile is not supported on python 3.12+"
    )
    @unittest.skipIf(IS_WINDOWS, reason="Windows not yet supported for torch.compile")
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    @skipCUDAIfRocm
    def test_compile_with_dynamic_min_seq_len(self, device, dtype):
        # shape (B, *, D)
        # min seq len: 7
        nt = torch.nested.nested_tensor(
            [
                torch.randn(7, 5),
                torch.randn(8, 5),
                torch.randn(9, 5),
            ],
            layout=torch.jagged,
        )

        # min seq len: 8
        nt2 = torch.nested.nested_tensor(
            [
                torch.randn(8, 5),
                torch.randn(9, 5),
                torch.randn(10, 5),
            ],
            layout=torch.jagged,
        )

        def f(nt):
            # TODO: Replace with public API when we can use @properties
            return torch.ones_like(nt) * nt._get_min_seqlen()

        for dynamic in [False, True, None]:
            self.assertFalse(_recompiles_for_inputs(f, (nt,), (nt2,), dynamic=dynamic))

    @dtypes(torch.float32)
    @skipIfTorchDynamo("Test compiles internally")
    @unittest.skipIf(
        sys.version_info >= (3, 12), "torch.compile is not supported on python 3.12+"
    )
    @unittest.skipIf(IS_WINDOWS, reason="Windows not yet supported for torch.compile")
    @skipCUDAIf(not SM70OrLater, "GPU capability is < SM70")
    @skipCUDAIfRocm
    def test_compile_with_propagated_dynamic_max_seq_len(self, device, dtype):
        # shape (B, *, D)
        # max seq len: 18
        nt = torch.nested.nested_tensor(
            [
                torch.randn(2, 5),
                torch.randn(3, 5),
                torch.randn(18, 5),
            ],
            layout=torch.jagged,
        )

        # max seq len: 19
        nt2 = torch.nested.nested_tensor(
            [
                torch.randn(2, 5),
                torch.randn(3, 5),
                torch.randn(19, 5),
            ],
            layout=torch.jagged,
        )

        def f(nt):
            nt2 = nt.sin() + 1
            # TODO: Replace with public API when we can use @properties
            return torch.ones_like(nt2) * nt2._get_max_seqlen()

        ref = f(nt)
        output = torch.compile(f, fullgraph=True, dynamic=False)(nt)
        self.assertEqual(ref, output)

        for dynamic in [False, True, None]:
            self.assertFalse(_recompiles_for_inputs(f, (nt,), (nt2,), dynamic=dynamic))

    @dtypes(torch.float32, torch.double, torch.half)
    def test_unbind_backward(self, device, dtype):
        nt = torch.nested.nested_tensor(
            [
                torch.randn(2, 4, device=device),
                torch.randn(5, 4, device=device),
                torch.randn(3, 4, device=device),
            ],
            layout=torch.jagged,
            requires_grad=True,
        )

        a, b, c = nt.unbind()
        b.sum().backward()

        @torch._dynamo.disable
        def check(nt):
            expected_grad = torch.zeros_like(nt)
            expected_grad.unbind()[1].add_(1.0)
            self.assertEqual(nt.grad, expected_grad)

        check(nt)


FORWARD_FAILURES = {
    # === BEGIN NotImplementedError SECTION ===
    # unary
    "nn.functional.celu",
    "nn.functional.elu",
    "nn.functional.hardshrink",
    "nn.functional.hardsigmoid",
    "nn.functional.hardtanh",
    "nn.functional.logsigmoid",
    "nn.functional.mish",
    "nn.functional.relu6",
    "nn.functional.rrelu",
    "nn.functional.selu",
    "nn.functional.softplus",
    "nn.functional.softshrink",
    "nn.functional.threshold",
    "rad2deg",
    # binary
    "__rsub__",
    "complex",
    "floor_divide",
    "polar",
    "rsub",
    # reduction
    "all",
    "amax",
    "amin",
    "any",
    "argmax",
    "argmin",
    "count_nonzero",
    "linalg.vector_norm",
    "nansum",
    "std",
    "std.unbiased",
    "var",
    "var.unbiased",
    # === BEGIN UNSUPPORTED SECTION ===
    # RuntimeError: mean(): not supported for NestedTensor on dim=1
    "mean",
    # ValueError: expects strided tensor (got torch.jagged tensor)
    "masked.amax",
    "masked.amin",
    "masked.argmax",
    "masked.argmin",
    "masked.logsumexp",
    "masked.mean",
    "masked.norm",
    "masked.prod",
    "masked.std",
    "masked.sum",
    "masked.var",
    # === BEGIN BUG SECTION ===
    # Returns a tuple of Tensors so it doesn't work with NJT's unary pointwise logic
    "frexp",
    # Need to adjust sample input func to pass the right thing
    "nn.functional.prelu",
    # TypeError: fill() received an invalid combination of arguments
    # got (NestedTensor), but expected one of:
    # * (Tensor input, Tensor value)
    # * (Tensor input, Number value)
    "fill",
    # RuntimeError: unsupported tensor layout: Jagged
    "jiterator_binary",
    "jiterator_binary_return_by_ref",
    "jiterator_unary",
    # Bug found: sum() with keepdim=True returns invalid shape
    "sum",
    # RuntimeError: prod(): keepdim=True must be set for NestedTensor
    "prod",
    # RuntimeError: "jagged_to_padded_dense" not implemented for 'Bool'
    "nanmean",
}

BACKWARD_FAILURES = {
    *FORWARD_FAILURES,
    # TODO: categorize these
    "__rpow__",
    "atanh",
    "cdouble",
    "cfloat",
    "chalf",
    "clamp_max",
    "clamp_min",
    "copysign",
    "float_power",
    "max.binary",
    "maximum",
    "min.binary",
    "minimum",
    "pow",
    "sgn",
    "sinc",
    "special.i1",
    "special.i1e",
    # clone() on a "non-contiguous with holes" NJT allocates a new offsets -> new nested int
    # RuntimeError: Function CloneBackward0 returned an invalid gradient at index 0 -
    # got [3, j29, 5] but expected shape compatible with [3, j28, 5]
    "clone",
    # Calling into torch.ops.aten.size directly
    "masked_select",
}

COMPILE_FORWARD_FAILURES = {
    *FORWARD_FAILURES,
    # clone() on non-contiguous with holes NJTs currently use unbind(), leading to
    # data-dependent error in torch.compile
    "clone",
}

COMPARE_TENSOR_COMPONENT_EQUALITY = {
    # masked_select is expected to output a different shape
    "masked_select",
}


def withXFails(failure_list):
    return decorateIf(
        unittest.expectedFailure,
        lambda params: params["op"].full_name in failure_list,
    )


# OpInfo-based NJT tests. These tests utilize an NJT-specific op_db generated from the standard
# op_db. Note that certain tradeoffs were made wrt coverage vs. time spent running tests:
#   * All tests run with dtype=torch.float32 only
class TestNestedTensorOpInfo(NestedTensorTestCase):
    # TODO: move this
    def _gen_grad_outputs(self, out_val):
        if isinstance(out_val, (list, tuple)):
            return tuple(torch.ones_like(c) for c in out_val)
        else:
            return (torch.ones_like(out_val),)

    @withXFails(FORWARD_FAILURES)
    @ops([op for op in njt_op_db if op.supports_njt], allowed_dtypes=(torch.float32,))
    def test_forward(self, device, dtype, op):
        for sample in op.sample_inputs(device=device, dtype=dtype, requires_grad=False):
            # compare to reference, but expect different nested int
            out = op.op(sample.input, *sample.args, **sample.kwargs)
            out_ref = op.ref(op, sample)
            self.assertEqualIgnoringNestedInts(out, out_ref)

    @withXFails(BACKWARD_FAILURES)
    @ops(
        [op for op in njt_op_db if op.supports_njt and op.supports_autograd],
        allowed_dtypes=(torch.float32,),
    )
    def test_backward(self, device, dtype, op):
        for sample in op.sample_inputs(device=device, dtype=dtype, requires_grad=True):
            # compare to reference, but expect different nested int
            out = op.op(sample.input, *sample.args, **sample.kwargs)
            out_ref = op.ref(op, sample)
            self.assertEqualIgnoringNestedInts(out, out_ref)

            inps, _ = tree_flatten((sample.input, sample.args, sample.kwargs))
            g_inps = [
                inp
                for inp in inps
                if isinstance(inp, torch.Tensor) and inp.requires_grad
            ]
            if len(g_inps) > 0:
                grads = torch.autograd.grad(
                    out, inputs=g_inps, grad_outputs=self._gen_grad_outputs(out)
                )

                grads_ref = torch.autograd.grad(
                    out_ref,
                    inputs=g_inps,
                    grad_outputs=self._gen_grad_outputs(out_ref),
                )

                self.assertEqual(grads, grads_ref)

    @withXFails(COMPILE_FORWARD_FAILURES)
    @torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True)
    @ops([op for op in njt_op_db if op.supports_njt], allowed_dtypes=(torch.float32,))
    def test_compile_forward(self, device, dtype, op):
        for sample in op.sample_inputs(device=device, dtype=dtype, requires_grad=False):
            torch.compiler.reset()

            op_fn = op.op

            def f(*args, **kwargs):
                return op_fn(*args, **kwargs)

            compiled_f = torch.compile(
                f, fullgraph=True, backend="aot_eager_decomp_partition"
            )

            out_ref = f(sample.input, *sample.args, **sample.kwargs)
            out_compile = compiled_f(sample.input, *sample.args, **sample.kwargs)

            if op.full_name in COMPARE_TENSOR_COMPONENT_EQUALITY:
                self.assertEqualIgnoringNestedInts(out_compile, out_ref)
            else:
                self.assertEqual(out_compile, out_ref)

    @withXFails(BACKWARD_FAILURES)
    @ops(
        [op for op in njt_op_db if op.supports_njt and op.supports_autograd],
        allowed_dtypes=(torch.float32,),
    )
    @torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True)
    def test_compile_backward(self, device, dtype, op):
        for sample in op.sample_inputs(device=device, dtype=dtype, requires_grad=True):
            torch.compiler.reset()

            op_fn = op.op

            def f(*args, **kwargs):
                return op_fn(*args, **kwargs)

            compiled_f = torch.compile(
                f, fullgraph=True, backend="aot_eager_decomp_partition"
            )

            out_ref = f(sample.input, *sample.args, **sample.kwargs)
            out_compile = compiled_f(sample.input, *sample.args, **sample.kwargs)

            self.assertEqual(out_compile, out_ref)

            inps, _ = tree_flatten((sample.input, sample.args, sample.kwargs))
            g_inps = [
                inp
                for inp in inps
                if isinstance(inp, torch.Tensor) and inp.requires_grad
            ]
            if len(g_inps) > 0:
                grads_compile = torch.autograd.grad(
                    out_compile,
                    inputs=g_inps,
                    grad_outputs=self._gen_grad_outputs(out_compile),
                )

                grads_ref = torch.autograd.grad(
                    out_ref, inputs=g_inps, grad_outputs=self._gen_grad_outputs(out_ref)
                )

                self.assertEqual(grads_compile, grads_ref)


instantiate_parametrized_tests(TestNestedTensor)
instantiate_device_type_tests(TestNestedTensorDeviceType, globals())
instantiate_device_type_tests(TestNestedTensorAutograd, globals())
instantiate_device_type_tests(TestNestedTensorSubclass, globals())
instantiate_device_type_tests(TestNestedTensorOpInfo, globals())

if __name__ == "__main__":
    run_tests()