Add torch.stack support

[yf225]

Fixes #523.

[jansel]

Could this be handled similar to:

helion/helion/_compiler/inductor_lowering.py

Lines 823 to 836 in 3c0348a

	@register_lowering(
	torch.ops.aten.permute.default, # pyright: ignore[reportAttributeAccessIssue]
	masked_value_fn=passthrough_masked_value,
	)
	def codegen_permute(ctx: GraphInterpreter, node: torch.fx.Node) -> object:
	assert not node.kwargs, "getitem kwargs not supported"
	tensor, dims = map_arg(node.args, lambda arg: ctx.env[arg])
	assert isinstance(tensor, ast.AST)
	dims = [*dims] # pyright: ignore[reportGeneralTypeIssues,reportOptionalIterable]
	assert {*dims} == {*range(len(dims))}, dims
	return expr_from_string(
	f"tl.permute({{tensor}}, {dims!r})",
	tensor=tensor,
	)

I'd expect that to be simpler with less need for special casing view ops.

[yf225]

Normally, torch.stack is decomposed to unsqueeze + cat, but I haven’t figure out a way to make codegen_cat work, so as a workaround we disable the decomp for torch.stack and implement codegen_stack instead.

[yf225]

Added test to make sure _get_custom_decomp_table doesn't affect normal torch.compile decomp for torch.stack