Indices should be either on cpu or on the same device as the indexed tensor

aiden-leong: <pre><code class="lang-auto">Using device: mps
1.13.0.dev20220614
0.14.0.dev20220614
</code></pre>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "Disco_Diffusion_v5_2_m1.py", line 2340, in &lt;module&gt;
    do_run()
  File "Disco_Diffusion_v5_2_m1.py", line 983, in do_run
    txt = clip_model.encode_text(clip.tokenize(prompt).to(device)).float()
  File "/Users/aiden/notebook/CLIP/clip/model.py", line 355, in encode_text
    x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
</code></pre>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/7b3a0ff87a91c98574ddfd98532062d78c49b2c7/aten/src/ATen/native/TensorAdvancedIndexing.cpp#L371">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/7b3a0ff87a91c98574ddfd98532062d78c49b2c7/aten/src/ATen/native/TensorAdvancedIndexing.cpp#L371" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/7b3a0ff87a91c98574ddfd98532062d78c49b2c7/aten/src/ATen/native/TensorAdvancedIndexing.cpp#L371" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/7b3a0ff87a91c98574ddfd98532062d78c49b2c7/aten/src/ATen/native/TensorAdvancedIndexing.cpp#L371</a></h4>



    <pre class="onebox">      <code class="lang-cpp">
        <ol class="start lines" start="361" style="counter-reset: li-counter 360 ;">
            <li>    const Tensor&amp; self,</li>
            <li>    const at::MaterializedIOptTensorListRef&amp; indices) {</li>
            <li>  auto dev = self.device();</li>
            <li>  bool indices_on_cpu_or_dev = std::all_of(</li>
            <li>      indices.begin(), indices.end(), [=](const at::OptionalTensorRef&amp; opt) {</li>
            <li>        return opt.has_value() ? (opt-&gt;is_cpu() || opt-&gt;device() == dev) : true;</li>
            <li>      });</li>
            <li>  TORCH_CHECK(</li>
            <li>      indices_on_cpu_or_dev,</li>
            <li>      "indices should be either on ", kCPU,</li>
            <li class="selected">      " or on the same device as the indexed tensor (", dev, ")");</li>
            <li>}</li>
            <li>
            </li>
<li>TORCH_PRECOMPUTE_META_FUNC2(index, Tensor)</li>
            <li>(const Tensor&amp; self, at::IOptTensorListRef indices) {</li>
            <li>  auto materialized = indices.materialize();</li>
            <li>
            </li>
<li>  TORCH_CHECK_INDEX(</li>
            <li>      materialized.size() &lt;= (size_t)self.dim(),</li>
            <li>      "too many indices for tensor of dimension ",</li>
            <li>      self.dim(), " (got ", materialized.size(), ")");</li>
        </ol>
      </code>
    </pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


albanD: <p>Can you make sure that all the Tensors are on the same device?<br>
In particular, you might want to use <code>torch.arange(size, device=x.device)</code> to make sure it matches.</p>

aiden-leong: <aside class="quote no-group" data-username="albanD" data-post="2" data-topic="154141">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>, device=x.device</p>
</blockquote>
</aside>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/a/dad2fcddaea9e305662a5b6424b8c728ce0b951d.png" data-download-href="https://discuss.pytorch.org/uploads/default/dad2fcddaea9e305662a5b6424b8c728ce0b951d" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/dad2fcddaea9e305662a5b6424b8c728ce0b951d_2_690x440.png" alt="image" data-base62-sha1="vdO7V2UXcaw5A1iwirghueP7fQh" width="690" height="440" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/dad2fcddaea9e305662a5b6424b8c728ce0b951d_2_690x440.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/dad2fcddaea9e305662a5b6424b8c728ce0b951d_2_1035x660.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/dad2fcddaea9e305662a5b6424b8c728ce0b951d_2_1380x880.png 2x" data-dominant-color="F2D2CD"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1690×1080 141 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<pre><code class="lang-auto">        print("+++++++++++++++")
        print(x.device)
        print(text.device)
        print(self.text_projection.device)
        print("+++++++++++++++")
</code></pre>
<p>They are all on device mps:0</p>

aiden-leong: <pre><code class="lang-auto">import numpy as np
x = torch.arange(72, dtype=torch.float32).to(device)
x = x.reshape((6,3,4))
print(x)
print(x.device)
index1 = torch.arange(1).to(device)
index2 = torch.arange(1).to(device)
index2[0] = 5
print(index1)
print(index2)
y = x[index1, index2]
print(y)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/b/db9ee110111b21e496f448fcbaaf2d224593b26a.png" data-download-href="https://discuss.pytorch.org/uploads/default/db9ee110111b21e496f448fcbaaf2d224593b26a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/b/db9ee110111b21e496f448fcbaaf2d224593b26a_2_690x175.png" alt="image" data-base62-sha1="vkQXHVtIpKja7NEduSVanXKCuEO" width="690" height="175" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/b/db9ee110111b21e496f448fcbaaf2d224593b26a_2_690x175.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/b/db9ee110111b21e496f448fcbaaf2d224593b26a_2_1035x262.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/b/db9ee110111b21e496f448fcbaaf2d224593b26a_2_1380x350.png 2x" data-dominant-color="EEC6C7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1916×486 89 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<pre><code class="lang-auto">NotImplementedError: The operator 'aten::index.Tensor_out' is not current implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.
</code></pre>
<p>‘aten::index.Tensor_out’ triggers fallback to cpu.</p>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/77764#issuecomment-1152966536">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/77764#issuecomment-1152966536" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/77764#issuecomment-1152966536" target="_blank" rel="noopener nofollow ugc">General MPS op coverage tracking issue</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2022-05-18" data-time="18:12:47" data-timezone="UTC">06:12PM - 18 May 22 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/albanD" target="_blank" rel="noopener nofollow ugc">
          <img alt="albanD" src="https://avatars.githubusercontent.com/u/6359743?v=4" class="onebox-avatar-inline" width="20" height="20">
          albanD
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          feature
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triaged
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: mps
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">### This issue is to have a centralized place to list and track work on adding s<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">upport to new ops for the MPS backend.

[**MPS operators coverage matrix**](https://qqaatw.github.io/pytorch-mps-ops-coverage/) - The matrix covers most of the supported operators but is not exhaustive. Before you comment below, please take a look at this matrix to make sure the operator you're requesting has not been implemented in nightly.  More details can be found on the [readme](https://github.com/qqaatw/pytorch-mps-ops-coverage).

There are a very large number of operators in pytorch and so they are not all implemented yet for the MPS backends as it is still in the prototype phase. We will be prioritizing adding new operators based on user feedback. If possible, please also provide link to the network or use-case where this op is getting used.

If you want to work on adding support for such op, feel free to comment below to get assigned one. Please avoid pickup up an op that is already being worked on or that already has a PR associated with it.

[Link to the wiki for details](https://github.com/pytorch/pytorch/wiki/MPS-Backend) on how to add these ops and example PRs.

**Good First Issue:**
Below is list of Ops which are good to get started to add operations to MPS backend. Please consider picking them up. 
- [ ] `nn.Conv3D`
- [ ] `aten::_weight_norm_interface`
- [ ] `aten::max_unpool2d`
- [ ] `aten::cummin.out`, `aten::cummax.out`
- [ ] `aten::upsample_linear1d.out`
- [ ] `aten::lerp.Scalar_out`
- [ ] `aten::renorm`


**Not categorized:**
These are the ops which are not yet picked up and need MPS implementation.

- [ ] `aten::slow_conv3d_forward` 
- [ ] `aten::_ctc_loss`
- [ ] `aten::avg_pool3d.out`
- [ ] `aten::linalg_qr.out`
- [ ] `aten::multilabel_margin_loss_forward`
- [ ] `aten::unique_dim`
- [ ] `aten::_sample_dirichlet`
- [ ] `aten::_fft_r2c`
- [ ] `aten::upsample_bicubic2d.out`
- [ ] `aten::linalg_inv_out_helper`
- [ ] `aten::bucketize`
- [ ] `aten::_embedding_bag`
- [ ] `aten::_standard_gamma`
- [ ] `aten::_upsample_bicubic2d_aa.out`
- [ ] `aten::'aten::_symeig_helper`
- [ ] `aten::linalg_matrix_exp`
- [ ] `aten::_nested_tensor_from_mask`
- [ ] `aten::randperm.generator_out`
- [ ] `aten::_fused_sdp_choice`
- [ ] `aten::linalg_cholesky_ex`
- [ ] `aten::scatter_reduce.two_out`
- [ ] `aten::kthvalue.values`
- [ ] `aten::_linalg_solve_ex.result`
- [ ] `aten::grid_sampler_2d_backward'`

**WIP:**
- [ ] `max_pool3d` https://github.com/pytorch/pytorch/pull/102148
- [ ] `aten::kl_div_backward` (Is not needed )


**Implemented Ops:**
Ops that have MPS backend implementations.

See [**MPS operators coverage matrix**](https://qqaatw.github.io/pytorch-mps-ops-coverage/) and the [readme](https://github.com/qqaatw/pytorch-mps-ops-coverage) for more details.

&lt;details&gt;
&lt;summary&gt;deprecated list&lt;/summary&gt;

- [x] `aten::histc` #96652
- [x] `pow.Scalar_out` (@qqaatw )
- [x] `aten::log_sigmoid_forward` (@qqaatw )
- [x] `aten::fmax.out` (@qqaatw )
- [x] `aten::roll`  https://github.com/pytorch/pytorch/pull/95168
- [x] `aten::hardsigmoid` (@qqaatw )
- [x] `aten::logit` (@qqaatw )
- [x] `linalg_solve_triangular`
- [x] `aten::sort.values_stable` https://github.com/pytorch/pytorch/issues/86750
- [x] `aten::remainder.Tensor_out` https://github.com/pytorch/pytorch/issues/86806
- [x] `aten::hardswish` https://github.com/pytorch/pytorch/issues/86807
- [x] `aten::nansum` https://github.com/pytorch/pytorch/issues/86809
- [x] `aten::fmod.Tensor_out` https://github.com/pytorch/pytorch/issues/86810
- [x] `aten::range` https://github.com/pytorch/pytorch/issues/86990
- [x] `aten::argsort` https://github.com/pytorch/pytorch/issues/86991
- [x] `aten::repeat_interleave` https://github.com/pytorch/pytorch/issues/87219
- [x] `aten::median` https://github.com/pytorch/pytorch/issues/87220
- [x] `aten::trace` https://github.com/pytorch/pytorch/issues/87221
- [x] `aten::im2col` (Falling back to CPU as its mostly used in preprocessing layers)
- [x] `aten::_cdist_forward` https://github.com/pytorch/pytorch/pull/91643
- [x] `aten::native_group_norm_backward` (Implemented by @malfet )
- [x] `aten::grid_sampler_2d` (https://github.com/pytorch/pytorch/pull/94273)
- [x] `aten::upsample_nearest1d_backward.grad_input`
- [x] `aten::upsample_nearest1d.out`
- [x] `aten::repeat_interleave.self_int`
- [x] `aten::nan_to_num.out`
- [x] `aten::unique_consecutive` https://github.com/pytorch/pytorch/pull/88532
- [x] `torch.bincount` https://github.com/pytorch/pytorch/pull/91267
- [x] `aten::_unique2`  https://github.com/pytorch/pytorch/pull/88532
- [x] `aten::unfold` https://github.com/pytorch/pytorch/pull/91266
- [x] `aten::triangular_solve.X` https://github.com/pytorch/pytorch/pull/94345
- [x] `aten::nonzero` https://github.com/pytorch/pytorch/pull/91616
- [x] `aten::_index_put_impl_` (https://github.com/pytorch/pytorch/pull/85672)
- [x] `aten::amax.out` (#79682)
- [X] `aten::_slow_conv2d_forward` (https://github.com/pytorch/pytorch/pull/86303)
- [x] `aten::eye.m_out` (https://github.com/pytorch/pytorch/pull/78408)
- [x] `aten::multinomial` (https://github.com/pytorch/pytorch/pull/80760 ) 
- [x] `aten::flip` (#80214)
- [x] `aten::equal` https://github.com/pytorch/pytorch/pull/80195
- [x] `aten::_local_scalar_dense`
- [x] `aten::l1_loss_backward.grad_input` (#80010)
- [x] `aten::glu.out` (#79866)
- [x] ` aten::linspace.out` https://github.com/pytorch/pytorch/pull/78570
- [x]  `aten::arange.out` https://github.com/pytorch/pytorch/pull/78789
- [x] `aten::adaptive_max_pool2d` https://github.com/pytorch/pytorch/pull/78410
- [x] `aten::count_nonzero.dim_IntList`
- [x] `aten::softplus.out` (https://github.com/pytorch/pytorch/pull/78930)
- [x] `aten::index_add.out` https://github.com/pytorch/pytorch/pull/79935
- [x] `aten::normal` (#80297)
- [x] `aten::native_layer_norm_backward` https://github.com/pytorch/pytorch/pull/79189
- [x] `aten::logical_and.out` (#80216)
- [x] `aten::frac.out` (https://github.com/pytorch/pytorch/pull/86625)
- [x] `aten:: masked_select` https://github.com/pytorch/pytorch/pull/85818
- [x] `aten::softplus_backward.grad_input` (#79873)
- [x] `aten::slow_conv_transpose2d.out` (@malfet could be due to incompatibility with torchvision)
- [x] `aten::signbit.out` (https://github.com/pytorch/pytorch/pull/87214)
- [X] `aten::cumsum.out` (https://github.com/pytorch/pytorch/pull/88319)
- [X] `aten::cumprod.out` 
- [X] `aten::expm1.out` (https://github.com/pytorch/pytorch/pull/87147)
- [x] `aten::bitwise_xor.Tensor_out` (https://github.com/pytorch/pytorch/pull/82307)
- [x] `aten::bitwise_and.Tensor_out` (https://github.com/pytorch/pytorch/pull/82307)
- [x] `aten::bitwise_or.Tensor_out` (https://github.com/pytorch/pytorch/pull/82307)
- [x] `aten::index.Tensor` (https://github.com/pytorch/pytorch/pull/82507)
- [x] `aten::index.Tensor_out` (https://github.com/pytorch/pytorch/pull/82507)
&lt;/details&gt;


**Ops not supported by MPS:**
Ops that will require either to use the CPU fallback system or a custom Metal kernel.

- [ ] `aten::lgamma.out`
- [ ] `aten::linalg_householder_product`</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>