Enable tensorrt.assertion in Tripy

[pranavm-nvidia]

This is slightly tricky because the assertion layer has no outputs, which won't work with our current tracing logic. One approach might be to start tracking all consumers of trace tensors and then use those to find Assert trace operations, which we can then include in the trace.

Once this is enabled, we should also add an assertion to the expand API to make sure we only broadcast singleton dimensions.