Rewrite dtype functions for softmax ops in Python

[gpetters94]

No description provided.

[gpetters94]

@ramiro050 these are all one-liners, do you think they require tests?

[ramiro050]

@ramiro050 these are all one-liners, do you think they require tests?

Yes, every dtype function should be tested. What's nice is that the testing code is also a one-liner (two-liner tops) for these ops 🙂

See below for examples of testing ops with one tensor inputs and two tensor inputs:

torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/abstract_interp_lib_gen.py

Lines 1305 to 1312 in d3a49fd

	@check_dtype_function(_check_tensors_with_the_same_dtype(num_of_tensors=1, other=0.0) +
	_check_tensors_with_the_same_dtype(num_of_tensors=1, other=0))
	def aten〇eq〇Scalar〡dtype(self_rank_dtype: Tuple[int, int], other: Union[int, float]) -> int:
	return torch.bool
	@check_dtype_function(_check_two_tensor_op())
	def aten〇eq〇Tensor〡dtype(self_rank_dtype: Tuple[int, int], other_rank_dtype: Tuple[int, int]) -> int:
	return torch.bool

[ramiro050]

The CI is failing on the refine-types lit tests. Since we are slowly removing logic from RefineTypes, you can just delete the failing MLIR test

[gpetters94]

@ramiro050 It's failing only on dynamo, is this a known issue? Here's the IR before failure:

func.func @forward(%arg0: !torch.vtensor<[3,2,4],f32>, %arg1: !torch.vtensor<[3,2,4],f32>) -> !torch.vtensor<[3,2,4],unk> {
  %int1 = torch.constant.int 1
  %none = torch.constant.none
  %int6 = torch.constant.int 6
  %0 = torch.aten.clone %arg0, %none : !torch.vtensor<[3,2,4],f32>, !torch.none -> !torch.vtensor<[3,2,4],f32>
  %1 = torch.aten.clone %arg1, %none : !torch.vtensor<[3,2,4],f32>, !torch.none -> !torch.vtensor<[3,2,4],f32>
  %2 = torch.aten._softmax_backward_data %0, %1, %int1, %int6 : !torch.vtensor<[3,2,4],f32>, !torch.vtensor<[3,2,4],f32>, !torch.int, !torch.int -> !torch.vtensor<[3,2,4],unk>
  return %2 : !torch.vtensor<[3,2,4],unk>
}

[ramiro050]

Can you use the helper functions for testing? These create invocations that make very thorough checks of the dtype functions. Same comment applies to the rest of the functions. See:

torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/abstract_interp_lib_gen.py

Lines 1305 to 1312 in d3a49fd

	@check_dtype_function(_check_tensors_with_the_same_dtype(num_of_tensors=1, other=0.0) +
	_check_tensors_with_the_same_dtype(num_of_tensors=1, other=0))
	def aten〇eq〇Scalar〡dtype(self_rank_dtype: Tuple[int, int], other: Union[int, float]) -> int:
	return torch.bool
	@check_dtype_function(_check_two_tensor_op())
	def aten〇eq〇Tensor〡dtype(self_rank_dtype: Tuple[int, int], other_rank_dtype: Tuple[int, int]) -> int:
	return torch.bool

[ramiro050]

There is already a patch that will handle the ops here: #1895. Can you undo these changes to avoid conflicts?

[ramiro050]

I think the reason the test is failing is because there is currently no support for indexing tuples in dtype functions. Can you change this to grad_output_rank, grad_output_dtype = grad_output_rank_dtype like the other functions here to see if it fixes it?

[ramiro050]

Hey George, any updates on the dtype functions?