transformer_engine: wrap checker_fn in langctx and cleanup (PR2473)

[kshitij12345]

This allows checker_fn to call TensorProxy methods.

Following snippet should work now:

import thunder
from thunder.extend import OperatorExecutor, register_executor
import torch

my_exec = OperatorExecutor("my_exec")
register_executor(my_exec)

def _my_checker(x, y):
    # Helper function as input shape can be (*, Hin)
    def _view_input_as_2d(x):
        shape = x.shape
        return x.view((-1, shape[-1]))

    _view_input_as_2d(x)
    return True

my_exec.register_implementation(
    thunder.torch.mul,
    checker=_my_checker,
    execution_transform=lambda x, y: x * y,
)

def foo(x, y):
    return x * y

cfunc = thunder.jit(foo, executors=[my_exec])

# Works when requires_grad=False!
# x = torch.randn(16, 16, requires_grad=False)
x = torch.randn(16, 16, requires_grad=True)
cfunc(x, x)

[t-vi]

I know it is late, but how about just decorating the checker function that has it with @langctx(Languages.TORCH) (from thunder.core.langctxs import langctx, Languages) for now (until we have figured out a bit more where we want the langctx)?
This is what we do for other checkers and transforms in executors, too, I would document that to be the way.

[carmocca]

This is useful. I remember having to do

    # cannot use `.view` so reduce the dimension manually
    # n_rows = X.view(-1, X.shape[-1]).shape[0]
    n_rows = reduce(int.__mul__, X.shape[:-1])

for the Unsloth RMSNorm kernel