Fix captured random state in compile

[angeloskath]

This is an interesting bug. In our attempts at thread-safety, compile now captures the random state of the thread that created the function and not the thread that called the function which can (and often is) different.

There isn't too clean of a solution that I could think of so please be as strict as possible in your judgment. The current approach is twofold

1. Make the random state a singleton object that fetches the thread local state. This makes it ok to capture on any thread. Accessing its one element will get you the correct array based on which thread you are on.
2. The ugly part of this fix is that tree_visit and tree_visit_update now have a special case where if they encounter a random state they instead operate on the underlying array. This is ugly because it is not necessarily consistent with tree_map or multi-tree tree_visit.

This closes ml-explore/mlx-lm#1444 and fixes ml-explore/mlx-lm#1439 .

[zcbenz]

I wonder if we can do this in a less intrusive way by replacing

mlx/python/src/transforms.cpp

Lines 560 to 563 in 4367c73

	// Replace tracers with originals in captured inputs
	if (!captured_inputs.is_none()) {
	tree_replace(captured_inputs, trace_captures, flat_in_captures);
	}

with

if (!captured_inputs.is_none()) { 
  trace_captures.push_back(random_state_sentinel());
  flat_in_captures.push_back(random_state_key());
  tree_replace(captured_inputs, trace_captures, flat_in_captures);
}

[angeloskath]

The problem is the random state can be in any place in the captured inputs. As a result we have two options

1. Walk the captured inputs and replace the random state with the thread local one before flattening
2. Do it when the random state is encountered while flattening and while writing back

Option 1 has the overhead of 2 extra tree walks which for say 100k nodes (like a big transformer or sth) can be non-negligible.

Option 2 is intrusive and ugly.

[zcbenz]

That makes sense, we can probably add a callback to tree_flatten to make it cheap to do replacement but it feels over-engineered and still does not fix the hook in tree_visit_update, I'm good with current approach and I can't come with a better one.