support python mlx.array creation from list of mlx.array's by Alex-X-W · Pull Request #325 · ml-explore/mlx

Alex-X-W · 2023-12-31T03:18:50Z

Proposed changes

Adds support for python mlx.array create from list of mlx.array's. Trying to tackle issue 287.

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
I have updated the necessary documentation (if needed)

Alex-X-W · 2023-12-31T03:20:30Z

-enum PyScalarT {
-  pybool = 0,
-  pyint = 1,
-  pyfloat = 2,
-  pycomplex = 3,
-};
-


Replace it with Dtype since we'd have to deal with all possible mlx dtypes when create from list of mlx.arrays. Not just python primitive types.

Alex-X-W · 2023-12-31T03:23:09Z


  Val val;
-  const uint8_t size;
+  uint8_t size;


Have to drop the constness here since we'd need the Dtype to be copy-assignable, which we are going to need in validate_shape() so that we have a Dtype defined first and assigned later to get the final promoted type.

Can you keep this const? See comment below.

Alex-X-W · 2023-12-31T03:24:46Z

      fill_vector(l.template cast<py::tuple>(), vals);
+    } else if (py::isinstance<array>(l)) {
+      auto arr = py::cast<array>(l);
+      arr.eval();


Need to call eval() here since arr may come from transformations so that it does not hold the actual raw data.. without the eval() then the data<T>() pointer would be uninitialized (actually zero) so that we would end up with segmentation fault.

awni · 2023-12-31T05:29:48Z

Thanks for adding this. I was thinking we should use concatenate and have this be part of the compute graph rather than force an eval.

Concretely when you do mx.array([a1, ...]), it's basically using concatenate under the hood. Does that make sense?

Alex-X-W · 2023-12-31T05:56:04Z

@awni Thanks. I see, yeah that makes sense. I am gonna renovate the code a bit.

Alex-X-W · 2024-01-01T08:42:26Z

+        throw std::invalid_argument(
+            "Initialization encountered non-uniform length.");
+      }


Tested in assertRaises here.

Alex-X-W · 2024-01-01T08:43:36Z

-  size_t size = 1;
-  for (auto s : shape) {
-    size *= s;


For some reason the size here is never used, so I deleted this block.

Oops, good catch.

Alex-X-W · 2024-01-01T08:45:28Z

Hi @awni ! May I bug you to review again? Thanks! And... Happy New Year 🥳

awni · 2024-01-02T18:24:03Z

Hi @Alex-X-W ,

There's two things I think we need to fix with this:

The main point to using stack is to avoid needing to eval the input arrays. I see that you eval them then get the raw pointer based on the type and copy it. That should change to be much simpler. Just put the arrays in a vector, the stack will handle type promotion.
If you are creating an array from a big list (or list of lists) of python scalars it will be really inefficient to convert these to arrays and do a stack. I would change the code so that when there is no nested array it follows a different path and puts the values in vector of the appropriate type and then uses the array constructor.

awni

Left a comment with requested changes.

…itialized by fill_vector

Alex-X-W · 2024-01-03T05:00:09Z

+
+  if (all_python_primitive_elements) {
+    // `pl` does not contain mlx arrays
+    return array_from_list(pl, type, dtype, shape);


This code path handles only the code path when all elements are python primitive types

Alex-X-W · 2024-01-03T05:01:26Z

+            mx.int64,
+            mx.float16,
+            mx.float32,
+            mx.bfloat16,


Added bfloat16 given #339 is merged.

awni

This is looking really nice!!!

My main comment is I would prefer not to undo the constness of the dtype.val unless there is a strong need and I don't think there is in this case (see inline comments). Could you make the change, then I think we are good to go?

awni · 2024-01-04T03:43:25Z


  Val val;
-  const uint8_t size;
+  uint8_t size;


Can you keep this const? See comment below.

awni · 2024-01-04T03:43:41Z

+    } else if (py::isinstance<array>(l)) {
+      all_python_primitive_elements = false;
+      auto arr = py::cast<array>(l);
+      if (arr.shape().size() + idx + 1 == shape.size() &&


nit: arr.ndim() instead of arr.shape().size()

awni · 2024-01-04T03:50:45Z

+  Dtype type = bool_;
  for (auto l : list) {
-    PyScalarT t;
+    Dtype t = bool_;


To avoid needing to change the const of the dtype value, you could revert this to use the PyScalarT and if you get an array, just set it bool_ (lowest order type) and let the stack handle the type promotion if you actually get any arrays (which it already will by default).

Alternatively you could put it in a lambda and add the types to a vector and the reduce the vector with promotion

I think the first option is simpler though.

…type promotions automatically

Alex-X-W · 2024-01-04T06:02:11Z

Thank you @awni for the pointers! They all make sense to me.
Please help take another look, thanks!

awni

Awesome!! I left one more nit. Could you check it then we can merge?

awni · 2024-01-04T15:53:48Z

      return get_shape(l.template cast<py::tuple>(), shape);
+    } else if (py::isinstance<array>(l)) {
+      auto arr = py::cast<array>(l);
+      shape.insert(shape.end(), arr.shape().begin(), arr.shape().end());


nit: I think shape = arr.shape() is sufficient here.

@awni no...
if a, b = mx.array([1]), mx.array([2])
then in evaluating mx.array([a, b]) on line 198 with shape = arr.shape() we'd end up with shape = [1], but actually the shape should be [2, 1]. shape = arr.shape() will drop a dimension in the process of list unrolling at line 189. You will see what I mean if you see the full block of the get_shape()

mlx/python/src/array.cpp

Lines 188 to 202 in 1bcee20

void get_shape(T list, std::vector<int>& shape) {

shape.push_back(py::len(list));

if (shape.back() > 0) {

auto& l = *list.begin();

if (py::isinstance<py::list>(l)) {

return get_shape(l.template cast<py::list>(), shape);

} else if (py::isinstance<py::tuple>(l)) {

return get_shape(l.template cast<py::tuple>(), shape);

} else if (py::isinstance<array>(l)) {

auto arr = py::cast<array>(l);

shape.insert(shape.end(), arr.shape().begin(), arr.shape().end());

return;

}

}

}

Sorry I misunderstood this code, thanks for clarifying.

awni · 2024-01-04T15:58:26Z

PS checked cpu/gpu tests pass + formatting passes

awni

🚀 Thanks so much for adding this!

limit the range in symmetric_matrices

Alex-X-W commented Dec 31, 2023

View reviewed changes

Alex-X-W force-pushed the mutexuan_array_from_arrays branch from 352c625 to ebe6ddd Compare January 1, 2024 08:36

Alex-X-W commented Jan 1, 2024

View reviewed changes

awni self-requested a review January 2, 2024 18:21

awni requested changes Jan 2, 2024

View reviewed changes

Alex-X-W added 3 commits January 3, 2024 10:15

support python mlx.array creation from list of mlx.array's

67b1146

include bfloat16 in UT

ab244d1

refactor so that sub array made of all python primitive types gets in…

97ad937

…itialized by fill_vector

Alex-X-W force-pushed the mutexuan_array_from_arrays branch from ebe6ddd to 97ad937 Compare January 3, 2024 04:54

Alex-X-W commented Jan 3, 2024

View reviewed changes

Alex-X-W requested a review from awni January 3, 2024 05:01

awni requested changes Jan 4, 2024

View reviewed changes

Alex-X-W added 2 commits January 4, 2024 13:11

address PR comment: arr.shape().size() -> arr.ndim()

fb84edc

address PR comment: get back Dtype constness and let stack to handle …

1bcee20

…type promotions automatically

Alex-X-W requested a review from awni January 4, 2024 06:02

awni approved these changes Jan 4, 2024

View reviewed changes

awni approved these changes Jan 5, 2024

View reviewed changes

awni merged commit d8f41a5 into ml-explore:main Jan 5, 2024

awni mentioned this pull request Jan 8, 2024

Support initializing array from list of arrays #287

Closed

zcbenz pushed a commit to Ritabanm/mlx that referenced this pull request Sep 30, 2026

Merge pull request ml-explore#325 from ev-br/eigvals_health_check

1572b12

limit the range in symmetric_matrices

	void get_shape(T list, std::vector<int>& shape) {
	shape.push_back(py::len(list));
	if (shape.back() > 0) {
	auto& l = *list.begin();
	if (py::isinstance<py::list>(l)) {
	return get_shape(l.template cast<py::list>(), shape);
	} else if (py::isinstance<py::tuple>(l)) {
	return get_shape(l.template cast<py::tuple>(), shape);
	} else if (py::isinstance<array>(l)) {
	auto arr = py::cast<array>(l);
	shape.insert(shape.end(), arr.shape().begin(), arr.shape().end());
	return;
	}
	}
	}

Conversation

Alex-X-W commented Dec 31, 2023

Proposed changes

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Alex-X-W Dec 31, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

awni commented Dec 31, 2023

Uh oh!

Alex-X-W commented Dec 31, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Alex-X-W commented Jan 1, 2024

Uh oh!

awni commented Jan 2, 2024

Uh oh!

awni left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

awni left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Alex-X-W commented Jan 4, 2024

Uh oh!

awni left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

awni commented Jan 4, 2024

Uh oh!

awni left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Alex-X-W Dec 31, 2023 •

edited

Loading