Fix compile with non standard types (#745)

* refactor tree utils

* fix compile + tree code refactor

* Add an extra test

* add a few missing activations to docs

* hash structure

* Encode the full argument structure

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

python/src/trees.h

@@ -0,0 +1,60 @@
+// Copyright © 2023-2024 Apple Inc.
+#pragma once
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "mlx/array.h"
+
+namespace py = pybind11;
+using namespace mlx::core;
+
+void tree_visit(py::object tree, std::function<void(py::handle)> visitor);
+
+py::object tree_map(
+    const std::vector<py::object>& trees,
+    std::function<py::object(const std::vector<py::object>&)> transform);
+
+py::object tree_map(
+    py::object tree,
+    std::function<py::object(py::handle)> transform);
+
+void tree_visit_update(
+    py::object tree,
+    std::function<py::object(py::handle)> visitor);
+
+/**
+ * Fill a pytree (recursive dict or list of dict or list) in place with the
+ * given arrays. */
+void tree_fill(py::object& tree, const std::vector<array>& values);
+
+/**
+ * Replace all the arrays from the src values with the dst values in the
+ * tree.
+ */
+void tree_replace(
+    py::object& tree,
+    const std::vector<array>& src,
+    const std::vector<array>& dst);
+
+/**
+ * Flatten a tree into a vector of arrays. If strict is true, then the
+ * function will throw if the tree contains a leaf which is not an array.
+ */
+std::vector<array> tree_flatten(py::object tree, bool strict = true);
+
+/**
+ * Unflatten a tree from a vector of arrays.
+ */
+py::object tree_unflatten(
+    py::object tree,
+    const std::vector<array>& values,
+    int index = 0);
+
+std::pair<std::vector<array>, py::object> tree_flatten_with_structure(
+    py::object tree,
+    bool strict = true);
+
+py::object tree_unflatten_from_structure(
+    py::object structure,
+    const std::vector<array>& values,
+    int index = 0);