<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<p dir="auto">I've found that there is some direction that hiding the communication time while computing einsum ops.<br>
Link: <a href="https://dl.acm.org/doi/pdf/10.1145/3567955.3567959" rel="nofollow">https://dl.acm.org/doi/pdf/10.1145/3567955.3567959</a><br>
I think it would be implemented somewhere since the affiliation of the paper is Google and used in TPU chips.<br>
Is there any reference implementation that I could follow?</p>
<p dir="auto">Also, I found this information while reading the following paper.<br>
Link: <a href="https://arxiv.org/abs/2211.05102" rel="nofollow">https://arxiv.org/abs/2211.05102</a><br>
If anyone who knows some reference implementations or materials that help to implement it, please let me know.</p>
<p dir="auto">Thank you for your support.</p>
<p dir="auto">Please:<br>
Revert <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="740316314" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/4867" data-hovercard-type="pull_request" data-hovercard-url="/google/jax/pull/4867/hovercard" href="https://github.com/google/jax/pull/4867">#4867</a> because it messes up TensorFlow, for example (ignore the fact the code is useless):</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; from jax.numpy import array
&gt;&gt;&gt; from tensorflow.keras import Sequential
&gt;&gt;&gt; Sequential().predict(array([0]))
[...]
TypeError: Using a non-tuple sequence for multidimensional indexing is not allowed; use `arr[array(seq)]` instead of `arr[seq]`. See https://github.com/google/jax/issues/4564 for more information.
[...]"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">import</span> <span class="pl-s1">array</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">from</span> <span class="pl-s1">tensorflow</span>.<span class="pl-s1">keras</span> <span class="pl-k">import</span> <span class="pl-v">Sequential</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-v">Sequential</span>().<span class="pl-en">predict</span>(<span class="pl-en">array</span>([<span class="pl-c1">0</span>]))
[...]
<span class="pl-v">TypeError</span>: <span class="pl-v">Using</span> <span class="pl-s1">a</span> <span class="pl-s1">non</span><span class="pl-c1">-</span><span class="pl-s1">tuple</span> <span class="pl-s1">sequence</span> <span class="pl-k">for</span> <span class="pl-s1">multidimensional</span> <span class="pl-s1">indexing</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-s1">allowed</span>; <span class="pl-s1">use</span> <span class="pl-s">`arr[array(seq)]`</span> <span class="pl-s1">instead</span> <span class="pl-s1">of</span> <span class="pl-s">`arr[seq]`</span>. <span class="pl-v">See</span> <span class="pl-s1">https</span>:<span class="pl-c1">//</span><span class="pl-s1">github</span>.<span class="pl-s1">com</span><span class="pl-c1">/</span><span class="pl-s1">google</span><span class="pl-c1">/</span><span class="pl-s1">jax</span><span class="pl-c1">/</span><span class="pl-s1">issues</span><span class="pl-c1">/</span><span class="pl-c1">4564</span> <span class="pl-k">for</span> <span class="pl-s1">more</span> <span class="pl-s1">information</span>.
[...]</pre></div>
<p dir="auto">But if you run it with NumPy:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; from numpy import array
&gt;&gt;&gt; from tensorflow.keras import Sequential
&gt;&gt;&gt; Sequential().predict(array([0]))
array([0])"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">from</span> <span class="pl-s1">numpy</span> <span class="pl-k">import</span> <span class="pl-s1">array</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">from</span> <span class="pl-s1">tensorflow</span>.<span class="pl-s1">keras</span> <span class="pl-k">import</span> <span class="pl-v">Sequential</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-v">Sequential</span>().<span class="pl-en">predict</span>(<span class="pl-en">array</span>([<span class="pl-c1">0</span>]))
<span class="pl-en">array</span>([<span class="pl-c1">0</span>])</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>