nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: score-spec/score-compose
Pull request #152 (2024-06-19 15:21:57)

feat: hide deprecated run command
The run command is still accessible  and can still be used in CI if needed, but it is not removed from the documented behavior.



Patch:
diff --git a/README.md b/README.md
index 3e4356e..2d7cb17 100644
--- a/README.md
+++ b/README.md
@@ -136,7 +136,7 @@ Global Flags:
   -v, --verbose count   Increase log verbosity and detail by specifying this flag one or more times
 ```
 
-**NOTE**: The `score-compose run` command still exists but should be considered deprecated as it does not support resource provisioning.
+**NOTE**: The `score-compose run` command still exists but is hidden and should be considered deprecated as it does not support resource provisioning.
 
 ## ![Get involved](docs/images/get-involved.svg) Get involved
 
diff --git a/internal/command/run.go b/internal/command/run.go
index 8c2a641..0545ca0 100644
--- a/internal/command/run.go
+++ b/internal/command/run.go
@@ -76,7 +76,8 @@ func init() {
 var runCmd = &cobra.Command{
 	Use:   "run [--file=score.yaml] [--output=compose.yaml]",
 	Args:  cobra.NoArgs,
-	Short: "Translate the SCORE file to docker-compose configuration",
+	Short: "(Deprecated) Translate the SCORE file to docker-compose configuration",
+	Hidden: true,
 	RunE:  run,
 	// don't print the errors - we print these ourselves in main()
 	SilenceErrors: true,


Test patch:
diff --git a/internal/command/root_test.go b/internal/command/root_test.go
index 49b2b81..ed48340 100644
--- a/internal/command/root_test.go
+++ b/internal/command/root_test.go
@@ -38,7 +38,6 @@ Available Commands:
   help        Help about any command
   init        Initialise a new score-compose project with local state directory and score file
   resources   Subcommands related to provisioned resources
-  run         Translate the SCORE file to docker-compose configuration
 
 Flags:
   -h, --help            help for score-compose
diff --git a/internal/command/run_test.go b/internal/command/run_test.go
index f72b920..c00c20a 100644
--- a/internal/command/run_test.go
+++ b/internal/command/run_test.go
@@ -69,7 +69,7 @@ func executeAndResetCommand(ctx context.Context, cmd *cobra.Command, args []stri
 func TestRunHelp(t *testing.T) {
 	stdout, stderr, err := executeAndResetCommand(context.Background(), rootCmd, []string{"run", "--help"})
 	assert.NoError(t, err)
-	assert.Equal(t, `Translate the SCORE file to docker-compose configuration
+	assert.Equal(t, `(Deprecated) Translate the SCORE file to docker-compose configuration
 
 Usage:
   score-compose run [--file=score.yaml] [--output=compose.yaml] [flags]
Repository: score-spec/score-compose
Pull request #161 (2024-07-08 16:02:05)

fix: ensure resource output template has a newline
See #160.

Ensure there's a newline so that output in a terminal makes sense and doesn't mess up printing.



Patch:
diff --git a/internal/command/resources.go b/internal/command/resources.go
index e0a0c98..48f0a9a 100644
--- a/internal/command/resources.go
+++ b/internal/command/resources.go
@@ -18,6 +18,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log/slog"
+	"strings"
 	"text/template"
 
 	"github.com/Masterminds/sprig/v3"
@@ -95,6 +96,10 @@ be returned as json.
 				case "yaml":
 					return yaml.NewEncoder(cmd.OutOrStdout()).Encode(outputs)
 				default:
+					// ensure there is a new line at the end if one is not already present
+					if !strings.HasSuffix(formatValue, "\n") {
+						formatValue += "\n"
+					}
 					prepared, err := template.New("").Funcs(sprig.FuncMap()).Parse(formatValue)
 					if err != nil {
 						return fmt.Errorf("failed to parse format template: %w", err)


Test patch:
diff --git a/internal/command/resources_test.go b/internal/command/resources_test.go
index 5a51cd8..e0a1a81 100644
--- a/internal/command/resources_test.go
+++ b/internal/command/resources_test.go
@@ -19,6 +19,7 @@ import (
 	"encoding/json"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -111,6 +112,7 @@ volume.default#example.vol
 		assert.NoError(t, err)
 		var out map[string]interface{}
 		assert.NoError(t, json.Unmarshal([]byte(stdout), &out))
+		assert.True(t, strings.HasSuffix(stdout, "\n"))
 	})
 
 	t.Run("format yaml", func(t *testing.T) {
@@ -118,11 +120,18 @@ volume.default#example.vol
 		assert.NoError(t, err)
 		var out map[string]interface{}
 		assert.NoError(t, yaml.Unmarshal([]byte(stdout), &out))
+		assert.True(t, strings.HasSuffix(stdout, "\n"))
 	})
 
 	t.Run("format template", func(t *testing.T) {
 		stdout, _, err := executeAndResetCommand(context.Background(), rootCmd, []string{"resources", "get-outputs", "volume.default#example.vol", "--format", `{{ . | len }}`})
 		assert.NoError(t, err)
-		assert.Equal(t, "2", stdout)
+		assert.Equal(t, "2\n", stdout)
+	})
+
+	t.Run("format template with newline", func(t *testing.T) {
+		stdout, _, err := executeAndResetCommand(context.Background(), rootCmd, []string{"resources", "get-outputs", "volume.default#example.vol", "--format", "{{ . | len }}\n"})
+		assert.NoError(t, err)
+		assert.Equal(t, "2\n", stdout)
 	})
 }
Repository: jax-ml/jax
Pull request #25597 (2024-12-19 00:04:03)

Account for tokens in `allow_spmd_sharding_propagation_to_parameters` and `allow_spmd_sharding_propagation_to_output` compile options
Account for tokens in `allow_spmd_sharding_propagation_to_parameters` and `allow_spmd_sharding_propagation_to_output` compile options



Patch:
diff --git a/jax/_src/interpreters/pxla.py b/jax/_src/interpreters/pxla.py
index a5cd193b5..97df2c4bd 100644
--- a/jax/_src/interpreters/pxla.py
+++ b/jax/_src/interpreters/pxla.py
@@ -2909,9 +2909,9 @@ class UnloadedMeshExecutable:
       da = _create_da_object(tuple(device_assignment))
     del device_assignment
 
-    allow_prop_to_inputs = tuple(isinstance(i, (UnspecifiedValue, AUTO))
-                                 for i in in_shardings)
-    allow_prop_to_outputs = tuple(
+    allow_prop_to_inputs = (False,) * len(ordered_effects) + tuple(
+        isinstance(i, (UnspecifiedValue, AUTO)) for i in in_shardings)
+    allow_prop_to_outputs = (False,) * len(ordered_effects) + tuple(
         isinstance(o, (UnspecifiedValue, AUTO)) or mlir.contains_unconstrained(o)
         for o in out_shardings)
 


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 30ec7fb8f..8a7c0b2e6 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4305,6 +4305,24 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertLen(traced.in_avals[0], 1)
     self.assertLen(traced.in_avals[1], 0)  # empty kwarg
 
+  def test_empty_io_callback_under_shard_map(self):
+    if config.use_shardy_partitioner.value:
+      self.skipTest("Shardy errors out on empty callbacks.")
+    mesh = jtu.create_mesh((4,), 'i')
+
+    def empty_callback(x):
+      return
+
+    def _f(x, y):
+      jax.experimental.io_callback(
+          empty_callback, (), x, ordered=True)
+      return x + y[..., jnp.newaxis]
+
+    f = jax.jit(shard_map(
+        _f, mesh, in_specs=(P(None, 'i'), P(None)),
+        out_specs=P(None, 'i')))
+    f(jnp.zeros((2, 16)), jnp.ones(2))
+
   def test_jit_trace_lower_and_compile(self):
     def f(x):
       return x * 2
Repository: jax-ml/jax
Pull request #25642 (2024-12-20 18:01:49)

Fix backwards attention test for larger head dims
This PR fixes the bwd attention tests for head_dim > 32 and reduces the tolerance level.


Patch:
diff --git a/jax/experimental/pallas/ops/gpu/attention.py b/jax/experimental/pallas/ops/gpu/attention.py
index 198340ec0..ba2a496e2 100644
--- a/jax/experimental/pallas/ops/gpu/attention.py
+++ b/jax/experimental/pallas/ops/gpu/attention.py
@@ -598,7 +598,9 @@ def mha_reference(
 ):
   q_seq_len = q.shape[1]
   kv_seq_len = k.shape[1]
-  logits = jnp.einsum('bqhc,bkhc->bhqk', q, k).astype(jnp.float32)
+  logits = jnp.einsum(
+      'bqhc,bkhc->bhqk', q, k, preferred_element_type=jnp.float32
+  )
   mask = None
   if segment_ids is not None:
     mask = jnp.expand_dims(segment_mask(segment_ids, segment_ids), 1)
@@ -608,5 +610,7 @@ def mha_reference(
     causal_mask = jnp.broadcast_to(causal_mask, logits.shape)
     mask = causal_mask if mask is None else jnp.logical_and(mask, causal_mask)
   logits = logits if mask is None else jnp.where(mask, logits, float("-inf"))
-  weights = jax.nn.softmax(logits * sm_scale).astype(q.dtype)
-  return jnp.einsum('bhqk,bkhc->bqhc', weights, v)
+  weights = jax.nn.softmax(logits * sm_scale)
+  return jnp.einsum(
+      'bhqk,bkhc->bqhc', weights, v, preferred_element_type=jnp.float32
+  )


Test patch:
diff --git a/tests/pallas/BUILD b/tests/pallas/BUILD
index 667a65a1c..fcc6ee048 100644
--- a/tests/pallas/BUILD
+++ b/tests/pallas/BUILD
@@ -482,7 +482,7 @@ jax_multiplatform_test(
         "gpu_a100_x32",
         "gpu_h100_x32",
     ],
-    shard_count = 6,
+    shard_count = 15,
     deps = [
         "//jax:pallas",
         "//jax:pallas_gpu",
diff --git a/tests/pallas/gpu_ops_test.py b/tests/pallas/gpu_ops_test.py
index e3cc7b46e..a2ace08b2 100644
--- a/tests/pallas/gpu_ops_test.py
+++ b/tests/pallas/gpu_ops_test.py
@@ -224,8 +224,8 @@ class FusedAttentionTest(PallasBaseTest):
   @jtu.sample_product(
       batch_size=(1, 2),
       seq_len=(128, 384),
-      num_heads=(1, 2, 4),
-      head_dim=(32,),
+      num_heads=(1, 2),
+      head_dim=(32, 64, 128,),
       causal=(True, False),
       use_segment_ids=(True, False),
   )
@@ -266,9 +266,9 @@ class FusedAttentionTest(PallasBaseTest):
     dq, dk, dv = jax.grad(f, argnums=(0, 1, 2))(q, k, v)
     dq_ref, dk_ref, dv_ref = jax.grad(f_ref, argnums=(0, 1, 2))(q, k, v)
     # TODO(sharadmv): Fix test.
-    np.testing.assert_allclose(dq, dq_ref, atol=0.14)
-    np.testing.assert_allclose(dk, dk_ref, atol=0.14)
-    np.testing.assert_allclose(dv, dv_ref, atol=0.05)
+    self.assertAllClose(dq, dq_ref, atol=5e-2)
+    self.assertAllClose(dk, dk_ref, atol=5e-2)
+    self.assertAllClose(dv, dv_ref, atol=1e-3)
 
 
 class FusedAttentionInterpretTest(FusedAttentionTest):
Repository: jax-ml/jax
Pull request #25713 (2025-01-03 00:12:12)

jax.debug.print: respect trace-time np.printoptions
This makes it so that you can do
```python
with jnp.printoptions(...):
  jax.debug.print(...)
```
and the printed values will respect the specified print options.


Patch:
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
index 8bde9c162..769f9e44d 100644
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -300,8 +300,9 @@ class _DebugPrintFormatChecker(string.Formatter):
 
 formatter = _DebugPrintFormatChecker()
 
-def _format_print_callback(fmt: str, *args, **kwargs):
-  sys.stdout.write(fmt.format(*args, **kwargs) + "\n")
+def _format_print_callback(fmt: str, np_printoptions, *args, **kwargs):
+  with np.printoptions(**np_printoptions):
+    sys.stdout.write(fmt.format(*args, **kwargs) + "\n")
 
 def debug_print(fmt: str, *args, ordered: bool = False, **kwargs) -> None:
   """Prints values and works in staged out JAX functions.
@@ -338,8 +339,8 @@ def debug_print(fmt: str, *args, ordered: bool = False, **kwargs) -> None:
   # Check that we provide the correct arguments to be formatted.
   formatter.format(fmt, *args, **kwargs)
 
-  debug_callback(functools.partial(_format_print_callback, fmt), *args,
-                 **kwargs, ordered=ordered)
+  debug_callback(functools.partial(_format_print_callback, fmt, np.get_printoptions()),
+                 *args, **kwargs, ordered=ordered)
 
 
 # Sharding visualization


Test patch:
diff --git a/tests/debugging_primitives_test.py b/tests/debugging_primitives_test.py
index 5532fdf03..edf07f1bd 100644
--- a/tests/debugging_primitives_test.py
+++ b/tests/debugging_primitives_test.py
@@ -219,6 +219,29 @@ class DebugPrintTest(jtu.JaxTestCase):
          [ 1  2  3  4  5  6  7  8  9 10 12 13 14]]
     """))
 
+  def test_debug_print_respects_numpy_printoptions(self):
+    def f(x):
+      with np.printoptions(precision=2, suppress=True):
+        jax.debug.print("{}", x)
+    x = np.array([1.2345, 2.3456, 1E-7])
+
+    # Default numpy print options:
+    with jtu.capture_stdout() as output:
+      jax.debug.print("{}", x)
+    self.assertEqual(output(), "[1.2345e+00 2.3456e+00 1.0000e-07]\n")
+
+    # Modified print options without JIT:
+    with jtu.capture_stdout() as output:
+      f(x)
+      jax.effects_barrier()
+    self.assertEqual(output(), "[1.23 2.35 0.  ]\n")
+
+    # Modified print options with JIT:
+    with jtu.capture_stdout() as output:
+      jax.jit(f)(x)
+      jax.effects_barrier()
+    self.assertEqual(output(), "[1.23 2.35 0.  ]\n")
+
 
 class DebugPrintTransformationTest(jtu.JaxTestCase):
Repository: jax-ml/jax
Pull request #25743 (2025-01-06 23:05:47)

Make inspect_array_sharding work without mesh context manager too.
Make inspect_array_sharding work without mesh context manager too.



Patch:
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
index 769f9e44d..9d3d3f85f 100644
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -390,6 +390,10 @@ def _inspect_sharding_lowering_rule(ctx: mlir.LoweringRuleContext, value, *,
     if devices is None:
       raise AssertionError(
           'Please file a bug at https://github.com/jax-ml/jax/issues')
+    am = axis_context.abstract_mesh
+    if am is not None:
+      mesh = mesh_lib.Mesh(np.array(devices).reshape(am.axis_sizes),
+                           am.axis_names)
   elif isinstance(axis_context, sharding_impls.SPMDAxisContext):
     devices = axis_context.mesh._flat_devices_tuple
   else:


Test patch:
diff --git a/tests/debugging_primitives_test.py b/tests/debugging_primitives_test.py
index edf07f1bd..6afb41645 100644
--- a/tests/debugging_primitives_test.py
+++ b/tests/debugging_primitives_test.py
@@ -1133,9 +1133,9 @@ class InspectShardingTest(jtu.JaxTestCase):
     f(np.arange(8, dtype=jnp.float32))
     self.assertTrue(is_called)
 
-  def test_inspect_sharding_3d_input_pos_sharding(self):
+  def test_inspect_sharding_3d_jit(self):
     def _cb(sd):
-      self.assertIsInstance(sd, jax.sharding.PositionalSharding)
+      self.assertIsInstance(sd, jax.sharding.NamedSharding)
       self.assertLen(sd.device_set, 2)
 
     def f_(x):
@@ -1149,7 +1149,7 @@ class InspectShardingTest(jtu.JaxTestCase):
 
     f(arr)
 
-  def test_inspect_sharding_3d_input_named_sharding(self):
+  def test_inspect_sharding_3d_pjit(self):
     def _cb(sd):
       self.assertIsInstance(sd, jax.sharding.NamedSharding)
       self.assertLen(sd.device_set, 2)
Repository: jax-ml/jax
Pull request #25781 (2025-01-08 18:34:47)

#sdy add repr for Sdy ArraySharding and DimSharding
#sdy add repr for Sdy ArraySharding and DimSharding



Patch:
diff --git a/jax/_src/sharding_impls.py b/jax/_src/sharding_impls.py
index e98957af2..4a7d0d9f4 100644
--- a/jax/_src/sharding_impls.py
+++ b/jax/_src/sharding_impls.py
@@ -125,6 +125,17 @@ class SdyDimSharding:
         is_closed=self.is_closed,
         priority=self.priority)
 
+  def __repr__(self):
+    return f'SdyDimSharding({self._custom_repr()})'
+
+  def _custom_repr(self):
+    axes_repr = ', '.join(f"'{a}'" for a in self.axes)
+    open_repr = ''
+    if not self.is_closed:
+      open_repr = ', ?' if self.axes else '?'
+    priority_repr = '' if self.priority is None else f'p{self.priority}'
+    return f'{{{axes_repr}{open_repr}}}{priority_repr}'
+
 
 @dataclasses.dataclass
 class SdyArraySharding:
@@ -146,6 +157,13 @@ class SdyArraySharding:
         mesh_attr,
         [dim_sharding.build() for dim_sharding in self.dimension_shardings])
 
+  def __repr__(self):
+    dim_sharding_repr = ', '.join(
+        d._custom_repr() for d in self.dimension_shardings)
+    device_id_repr = (f', device_ids={self.logical_device_ids}'
+                      if self.logical_device_ids is not None else '')
+    return f"SdyArraySharding([{dim_sharding_repr}]{device_id_repr})"
+
 
 @util.cache(max_size=4096, trace_context_in_key=False)
 def named_sharding_to_xla_hlo_sharding(


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index ca8f1deed..edad58278 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -6632,6 +6632,28 @@ class ShardyTest(jtu.JaxTestCase):
     lowered_str = jax.jit(f, in_shardings=[AUTO(mesh), AUTO(mesh)]).lower(x, x).as_text()
     self.assertIn('sdy.mesh @mesh = <["x"=8]>', lowered_str)
 
+  def test_array_sharding_repr_with_priority(self):
+    sharding = sharding_impls.SdyArraySharding(
+        mesh_shape=(('data', 4), ('model', 8), ('expert', 2)),
+        dimension_shardings=[
+            sharding_impls.SdyDimSharding(axes=['data', 'expert'], is_closed=True),
+            sharding_impls.SdyDimSharding(axes=['model'], is_closed=False, priority=2)])
+    self.assertEqual(repr(sharding), "SdyArraySharding([{'data', 'expert'}, {'model', ?}p2])")
+
+  def test_array_sharding_repr_with_logical_ids(self):
+    abstract_mesh = jax.sharding.AbstractMesh((('x', 4), ('y', 8), ('z', 2)))
+    ns = NamedSharding(abstract_mesh, P(('x', 'y'), 'z', P.UNCONSTRAINED, None),
+                       _logical_device_ids=[4, 5, 6, 7, 0, 1, 2, 3])
+    self.assertEqual(repr(ns._to_sdy_sharding(4)),
+                     "SdyArraySharding([{'x', 'y'}, {'z'}, {?}, {}], "
+                     "device_ids=[4, 5, 6, 7, 0, 1, 2, 3])")
+
+  def test_dimension_sharding_repr(self):
+    dim_sharding = sharding_impls.SdyDimSharding(
+        axes=['data', 'model'], is_closed=False, priority=2)
+    self.assertEqual(repr(dim_sharding),
+                     "SdyDimSharding({'data', 'model', ?}p2)")
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: jax-ml/jax
Pull request #25784 (2025-01-08 20:10:53)

Match the behavior on single host wrt multi-host if tiled=False. Fixes https://github.com/jax-ml/jax/issues/25783
Match the behavior on single host wrt multi-host if tiled=False. Fixes https://github.com/jax-ml/jax/issues/25783



Patch:
diff --git a/jax/experimental/multihost_utils.py b/jax/experimental/multihost_utils.py
index 79989583f..560b5c5cc 100644
--- a/jax/experimental/multihost_utils.py
+++ b/jax/experimental/multihost_utils.py
@@ -103,7 +103,8 @@ def _handle_array_process_allgather(inp, tiled):
   else:
     # All inputs here will be fully addressable.
     if jax.process_count() == 1:
-      return np.asarray(inp)
+      out = np.asarray(inp)
+      return np.expand_dims(out, axis=0) if not tiled else out
 
     devices = np.array(jax.devices()).reshape(jax.process_count(),
                                               jax.local_device_count())


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index 97bf71a52..afcdad376 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -700,10 +700,14 @@ class JaxArrayTest(jtu.JaxTestCase):
 
   def test_process_allgather_single_host(self):
     x = jnp.arange(8.)
-    out = multihost_utils.process_allgather(x)
+    out = multihost_utils.process_allgather(x, tiled=True)
     self.assertEqual(out.shape, x.shape)
     self.assertArraysEqual(out, x)
 
+    out = multihost_utils.process_allgather(x)
+    self.assertEqual(out.shape, (1, x.shape[0]))
+    self.assertArraysEqual(out, np.expand_dims(x, axis=0))
+
   @jtu.sample_product(
     dtype=jtu.dtypes.all,
     shape=[(), (10), (2, 3)],
Repository: jax-ml/jax
Pull request #25791 (2025-01-08 23:42:12)

fix grad(logsumexp) to produce 0s where `where` is False
Thanks @epiqueras for noticing and diagnosing this issue:

```python
x = jnp.array([0., 0. 100.])
where = jnp.array([True, True, False])
g = jax.grad(lambda x: lsp_special.logsumexp(x, where=where))(x)  # [0.5, 0.5, nan]
```

The fix here is just to project the input. (Alternatively we could write a custom_jvp.)


Patch:
diff --git a/jax/_src/ops/special.py b/jax/_src/ops/special.py
index 59ad594ef..76a249d36 100644
--- a/jax/_src/ops/special.py
+++ b/jax/_src/ops/special.py
@@ -69,6 +69,8 @@ def logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None,
     Either an array ``result`` or a pair of arrays ``(result, sign)``, depending
     on the value of the ``return_sign`` argument.
   """
+  if where is not None:
+    a = jnp.where(where, a, 0)
   if b is not None:
     a_arr, b_arr = promote_args_inexact("logsumexp", a, b)
     a_arr = jnp.where(b_arr != 0, a_arr, -jnp.inf)


Test patch:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
index c30e8da65..aa53d4a03 100644
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -202,6 +202,11 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
     y_actual = lsp_special.logsumexp(x, where=mask)
     self.assertAllClose(y_expected, y_actual, check_dtypes=False)
 
+  def testLogSumExpWhereGrad(self):
+    x = jnp.array([0., 0., 0., 0., 100.])
+    g = jax.grad(lambda x: lsp_special.logsumexp(x, where=jnp.arange(5) < 4))(x)
+    self.assertAllClose(g, jnp.array([0.25, 0.25, 0.25, 0.25, 0.]))
+
   @jtu.sample_product(
     shape=all_shapes,
     dtype=float_dtypes,