[JAX] Fixing importing in the encoder examples (#1600)

import te before te_jax
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

examples/jax/encoder/common.py

@@ -4,6 +4,7 @@
 """Shared functions for the encoder tests"""
 from functools import lru_cache
 
+import transformer_engine
 from transformer_engine_jax import get_device_compute_capability