Bug with as_strided_tensorimpl for MPS devices
Labels: triaged, module: mps

### 🐛 Describe the bug

I am experiencing a problem on my M1 Pro MacBook with training a Neural ODE model.

My use-case is quite extensive and depends on multiple python files and some custom integration routines, so I am unable to provide a trimmed-down version of my code. I have, however, provided a traceback of the error.

```python 

Traceback (most recent call last):
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/damage_neural/neural.py", line 370, in <module>
    r = ode.odeint_adjoint(model, y[0],t,block_size=time_chunk_size)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/pyoptmat/ode.py", line 656, in odeint_adjoint
    return wrapper.apply(solver, times, *adjoint_params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/autograd/function.py", line 551, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/pyoptmat/ode.py", line 583, in forward
    y = solver.integrate(times, cache_adjoint=True)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/pyoptmat/ode.py", line 370, in integrate
    result[k : k + self.n] = self.block_update(
                             ^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/pyoptmat/ode.py", line 497, in block_update
    dy = chunktime.newton_raphson_chunk(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/pyoptmat/chunktime.py", line 41, in newton_raphson_chunk
    R, J = fn(x)
           ^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/pyoptmat/ode.py", line 493, in RJ
    yd, yJ = func(times, y)
             ^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/damage_neural/neural.py", line 231, in forward
    dy_dot_dy = vmap(vmap(jacfwd(self.rate, argnums = 1)))(t, y, erate, T)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/apis.py", line 188, in wrapped
    return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 266, in vmap_impl
    return _flat_vmap(
           ^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 38, in fn
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 379, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/apis.py", line 188, in wrapped
    return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 266, in vmap_impl
    return _flat_vmap(
           ^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 38, in fn
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 379, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1132, in wrapper_fn
    results = vmap(push_jvp, randomness=randomness)(basis)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/apis.py", line 188, in wrapped
    return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 266, in vmap_impl
    return _flat_vmap(
           ^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 38, in fn
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 379, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1123, in push_jvp
    output = _jvp_with_argnums(func, args, basis, argnums=argnums, has_aux=has_aux)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 38, in fn
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 969, in _jvp_with_argnums
    result_duals = func(*duals)
                   ^^^^^^^^^^^^
  File "/Users/ganesh/ArgonneWork/BlackBox_NN/env/damage_neural/neural.py", line 193, in rate
    x = torch.cat([y, erate.unsqueeze(-1), T.unsqueeze(-1)], dim = -1)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: !self.is_mps() INTERNAL ASSERT FAILED at "/Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/TensorShape.cpp":1166, please report a bug to PyTorch. as_strided_tensorimpl does not work with MPS; call self.as_strided(...) instead
```

Please let me know if there are any further details you might require.

Thanks.

### Versions

ct_env
Collecting environment information...
PyTorch version: 2.2.0.dev20231018
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.5.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.0.40.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.4 (main, Jul  5 2023, 08:40:20) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.5.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] numpy==1.26.1
[pip3] torch==2.2.0.dev20231018
[conda] functorch                 2.0.0                    pypi_0    pypi
[conda] numpy                     1.25.2                   pypi_0    pypi
[conda] torch                     2.0.1                    pypi_0    pypi
(env) (base) ganesh@Ganeshs-MacBook-Pro-3 / % 



cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev