Add tests for export including control flow models and quantized models

[junpeiz]

Proposed changes

This PR adds more test cases for export/import models, including control flow models and quantized models.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• [N/A] I have updated the necessary documentation (if needed)

[awni]

Looks good, thanks!

[junpeiz]

The failed CI is about test_leaks, which seems not related to this PR.

FAIL [0.246s]: test_leaks (test_export_import.TestExportImport)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/distiller/project/python/tests/test_export_import.py", line 270, in test_leaks
    self.assertEqual(mem_pre, mem_post)
AssertionError: 1388564 != 20

@awni Do you happen to know what could cause this?

[awni]

You can rebase on this one once it lands and that test should pass: #2443

[junpeiz]

@awni Another unrelated test failed. To make sure it is not a transient error, I want to rerun the CI, but looks like I don't have permissions (as shown in the screenshot)

======================================================================
FAIL: test_pseudo_inverse (test_linalg.TestLinalg.test_pseudo_inverse)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/circleci/project/python/tests/test_linalg.py", line 238, in test_pseudo_inverse
    self.assertTrue(mx.allclose(A @ A_plus @ A, A))
AssertionError: array(False, dtype=bool) is not true

[awni]

@junpeiz the cuda failure related to the test you wrote is an issue. You can skip that test on the CUDA backend for now until we support quantization. Add it here.

The other failure for the pseudo inverse is a recent flaky test unrelated to this PR. We'll have to track that one down.

[junpeiz]

@awni All CIs are green now. Could you help to merge this PR? Thanks!

[awni]

Thanks for the contribution @junpeiz !