flax

096b8a02 - Check for leaves of type variablelib.Variable when getting sharding specs.