GitLab

* add swizzle in jax
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* added outer_impl
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* clean up FFI
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>