feat: support aten.atan2 converter

[chohk88]

Description

New feature to support aten.atan2 converter. I also add test case including edge case for inputs are zero.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

• New feature (non-breaking change which adds functionality)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[zewenli98]

Overall looks good to me. It seems that get_value_by_condition is a helper function. Can you move it to .../dynamo/conversion/converter_utils.py?

[gs-olive]

Does this function support other being a Python float or does it require TRTTensor?

[chohk88]

I think it requires TRTTensor because both input and output are Tensor

[gs-olive]

This can be removed - replace with usage of impl.condition.select, from here:

TensorRT/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py

Lines 102 to 113 in 76c9ebd

	def select(
	ctx: ConversionContext,
	target: Target,
	source_ir: Optional[SourceIR],
	name: str,
	input: TRTTensor,
	other: TRTTensor,
	condition: TRTTensor,
	) -> TRTTensor:
	select_layer = ctx.net.add_select(condition, input, other)
	set_layer_name(select_layer, target, name + "_select", source_ir)
	return select_layer.get_output(0)

[gs-olive]

Alternatively, you may need where:

TensorRT/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py

Lines 18 to 26 in 76c9ebd

	def where(
	ctx: ConversionContext,
	target: Target,
	source_ir: Optional[SourceIR],
	name: str,
	input: Union[TRTTensor, np.ndarray, torch.Tensor],
	other: Union[TRTTensor, np.ndarray, torch.Tensor],
	condition: Union[TRTTensor, np.ndarray, torch.Tensor],
	) -> TRTTensor:

[chohk88]

I removed get_value_by_condition and replaced it with impl.condition.select.

[gs-olive]

If you require one or both input tensors to be TRTTensor, consider using:

TensorRT/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

Lines 62 to 66 in 6b8ccbc

	@enforce_tensor_types(
	{
	0: (TRTTensor,),
	}
	)

[chohk88]

I add enforce_tensor_types for both input and other.

[chohk88]

Overall looks good to me. It seems that get_value_by_condition is a helper function. Can you move it to .../dynamo/conversion/converter_utils.py?

I removed get_value_by_condition and replaced it with impl.condition.select.

[chohk88]

@zewenli98 @gs-olive I've made the changes you suggested. I also noticed I missed the case when input or other is a single constant, so I fixed that using broadcast function. And, I added test cases to check it's all working right. Thanks for your helpful suggestions!

[narendasan]

LGTM