Revert "[fake tensor cache] Support index with non bool/int8 indices (pytorch#151477)"

This reverts commit bdb34f5. Reverted pytorch#151477 on behalf of https://github.com/wdvr due to reverting confusing ghstack state ([comment](pytorch#151477 (comment)))

Author: pytorchmergebot

test/test_fake_tensor.py

@@ -2158,31 +2158,6 @@ def test_cache_tuple_outputs(self):
  extract_tensor_metadata(b),
  )
 
-
- def test_cache_aten_index(self):
- with FakeTensorMode():
- x = torch.randn(4, 4, 4)
- idx_tensor1 = torch.tensor([0, 2, 3])
- idx_tensor2 = torch.tensor([0, 1, 2])
-
- FakeTensorMode.cache_clear()
- self.assertHitsMisses(0, 0)
-
- ref = torch.ops.aten.index(x, [None, idx_tensor1, idx_tensor2])
- self.assertHitsMisses(0, 3)
-
- res = torch.ops.aten.index(x, [None, idx_tensor1, idx_tensor2])
- self.assertHitsMisses(1, 3)
- self.assertEqual(extract_tensor_metadata(ref), extract_tensor_metadata(res))
-
- with FakeTensorMode():
- x = torch.randn(4, 4, 4)
- idx_tensor1 = torch.tensor([True, True, False, True])
- self.assertRaises(DynamicOutputShapeException, lambda: torch.ops.aten.index(x, [None, idx_tensor1]))
-
- idx_tensor1 = torch.tensor([1, -2, 3, -4], dtype=torch.int8)
- self.assertRaises(DynamicOutputShapeException, lambda: torch.ops.aten.index(x, [None, idx_tensor1]))
-
  @skipIfTorchDynamo("cache hit/miss changes with invoke_subgraph caching")
  def test_invoke_subgraph(self):
  """

torch/_subclasses/fake_tensor.py

@@ -1539,21 +1539,6 @@ def _validate_cache_key(
  raise _BypassDispatchCache("data dependent output")
 
  if torch.Tag.dynamic_output_shape in func.tags:
- if func is aten.index.Tensor:
- _, new_kwargs = normalize_function( # type: ignore[misc]
- func, args=args, kwargs=kwargs, normalize_to_only_use_kwargs=True # type: ignore[arg-type]
- )
- for index in new_kwargs["indices"]:
- # index calls nonzero for bool or int8 tensors, and
- # therefore has a dynamic shape output. For other dtypes,
- # the output shape depends on the input shape (and not data)
- if isinstance(index, torch.Tensor) and index.dtype in (
- torch.bool,
- torch.int8,
- ):
- raise _BypassDispatchCache("dynamic output shape")
- return
-
  raise _BypassDispatchCache("dynamic output shape")
 
  if torch.Tag.inplace_view in func.tags: