Merge pull request #1527 from rstudio/fixes

Fixes for JAX updates

Author: t-kalinowski

NAMESPACE

@@ -12,6 +12,7 @@ S3method("[[",python.builtin.super)
 S3method("[[",python_builtin_super_getter)
 S3method(Arg,keras.src.backend.Tensor)
 S3method(Arg,keras.src.backend.common.keras_tensor.KerasTensor)
+S3method(Ops,jax._src.export.shape_poly._DimExpr)
 S3method(Summary,keras_shape)
 S3method(all,equal.numpy.ndarray)
 S3method(as.array,jax.Array)
@@ -40,6 +41,7 @@ S3method(as.numeric,keras.src.backend.common.variables.KerasVariable)
 S3method(base::all.equal,keras.src.backend.Tensor)
 S3method(base::all.equal,keras.src.backend.common.keras_tensor.KerasTensor)
 S3method(base::all.equal,keras.src.backend.common.variables.KerasVariable)
+S3method(base::as.array,PIL.Image.Image)
 S3method(compile,keras.src.models.model.Model)
 S3method(destructure,keras_shape)
 S3method(evaluate,keras.src.models.model.Model)

NEWS.md

@@ -23,6 +23,8 @@
 - Added elastic deformation utilities for images: `layer_random_elastic_transform()`
   and the lower-level `op_image_elastic_transform()`.
 
+- Added `as.array()` support for `PIL.Image.Image` objects.
+
 - Transposed convolution utilities now follow the latest Keras API:
   `op_conv_transpose()` defaults `strides = 1` and the `layer_conv_*_transpose()`
   layers expose `output_padding` for precise shape control.
@@ -38,6 +40,11 @@
 
 - `layer_layer_normalization()` removes the `rms_scaling` argument.
 
+- Merging layers now capture `...` with tidy dots (fixes #1525).
+
+- Fixed Ops on JAX `_DimExpr` so symbolic shapes survive arithmetic with R
+  double scalars.
+
 - `layer_reshape()` can now accept `-1` as a sentinel for an automatically calculated axis size.
 
 - `layer_torch_module_wrapper()` gains an `output_shape` argument to help Keras

R/jax-methods.R

@@ -94,3 +94,20 @@ type_sum.keras.src.backend.jax.core.JaxVariable <- type_sum.keras.src.backend.ja
 
 # "keras.src.backend.Variable" too?
 # "keras.src.backend.common.variables.Variable" too?
+
+#' @exportS3Method Ops jax._src.export.shape_poly._DimExpr
+Ops.jax._src.export.shape_poly._DimExpr <- function(e1, e2) {
+  if (missing(e2)) {
+    return(e1)
+  }
+  conv <- function(x) {
+    if (is.double(x) && isTRUE(all(x == suppressWarnings(as.integer(x))))) {
+      storage.mode(x) <- "integer"
+    }
+    x
+  }
+  e1 <- conv(e1)
+  e2 <- conv(e2)
+  NextMethod()
+}
+

R/layers-merging.R

@@ -52,7 +52,7 @@ function (inputs, ...)
     args <- capture_args(list(input_shape = normalize_shape,
         batch_size = as_integer, batch_input_shape = normalize_shape),
         ignore = c("...", "inputs"))
-    dots <- split_dots_named_unnamed(list(...))
+    dots <- split_dots_named_unnamed(list2(...))
     if (missing(inputs))
         inputs <- NULL
     else if (!is.null(inputs) && !is.list(inputs))
@@ -116,7 +116,7 @@ function (inputs, ...)
     args <- capture_args(list(input_shape = normalize_shape,
         batch_size = as_integer, batch_input_shape = normalize_shape),
         ignore = c("...", "inputs"))
-    dots <- split_dots_named_unnamed(list(...))
+    dots <- split_dots_named_unnamed(list2(...))
     if (missing(inputs))
         inputs <- NULL
     else if (!is.null(inputs) && !is.list(inputs))
@@ -177,7 +177,7 @@ function (inputs, ..., axis = -1L)
     args <- capture_args(list(axis = as_axis, input_shape = normalize_shape,
         batch_size = as_integer, batch_input_shape = normalize_shape),
         ignore = c("...", "inputs"))
-    dots <- split_dots_named_unnamed(list(...))
+    dots <- split_dots_named_unnamed(list2(...))
     if (missing(inputs))
         inputs <- NULL
     else if (!is.null(inputs) && !is.list(inputs))
@@ -260,7 +260,7 @@ function (inputs, ..., axes, normalize = FALSE)
     args <- capture_args(list(axes = as_axis, input_shape = normalize_shape,
         batch_size = as_integer, batch_input_shape = normalize_shape),
         ignore = c("...", "inputs"))
-    dots <- split_dots_named_unnamed(list(...))
+    dots <- split_dots_named_unnamed(list2(...))
     if (missing(inputs))
         inputs <- NULL
     else if (!is.null(inputs) && !is.list(inputs))
@@ -322,7 +322,7 @@ function (inputs, ...)
     args <- capture_args(list(input_shape = normalize_shape,
         batch_size = as_integer, batch_input_shape = normalize_shape),
         ignore = c("...", "inputs"))
-    dots <- split_dots_named_unnamed(list(...))
+    dots <- split_dots_named_unnamed(list2(...))
     if (missing(inputs))
         inputs <- NULL
     else if (!is.null(inputs) && !is.list(inputs))
@@ -384,7 +384,7 @@ function (inputs, ...)
     args <- capture_args(list(input_shape = normalize_shape,
         batch_size = as_integer, batch_input_shape = normalize_shape),
         ignore = c("...", "inputs"))
-    dots <- split_dots_named_unnamed(list(...))
+    dots <- split_dots_named_unnamed(list2(...))
     if (missing(inputs))
         inputs <- NULL
     else if (!is.null(inputs) && !is.list(inputs))
@@ -446,7 +446,7 @@ function (inputs, ...)
     args <- capture_args(list(input_shape = normalize_shape,
         batch_size = as_integer, batch_input_shape = normalize_shape),
         ignore = c("...", "inputs"))
-    dots <- split_dots_named_unnamed(list(...))
+    dots <- split_dots_named_unnamed(list2(...))
     if (missing(inputs))
         inputs <- NULL
     else if (!is.null(inputs) && !is.list(inputs))
@@ -509,7 +509,7 @@ function (inputs, ...)
     args <- capture_args(list(input_shape = normalize_shape,
         batch_size = as_integer, batch_input_shape = normalize_shape),
         ignore = c("...", "inputs"))
-    dots <- split_dots_named_unnamed(list(...))
+    dots <- split_dots_named_unnamed(list2(...))
     if (missing(inputs))
         inputs <- NULL
     else if (!is.null(inputs) && !is.list(inputs))

R/r-utils.R

@@ -27,7 +27,9 @@ split_dots_named_unnamed <- function(dots) {
   if (is.null(nms))
     return(list(unnamed = dots, named = list()))
   named <- nzchar(nms)
-  list(unnamed = dots[!named], named = dots[named])
+  unnamed <- dots[!named]
+  names(unnamed) <- NULL
+  list(unnamed = unnamed, named = dots[named])
 }
 
 drop_nulls <- function(x, i = NULL) {

R/s3-methods.R

@@ -110,3 +110,7 @@ py_to_r__keras.src.utils.tracking.TrackedSet <- function(x) import("builtins")$l
 # }
 # rm(list = c("generic", "cls"))
 
+#' @exportS3Method base::as.array
+as.array.PIL.Image.Image <- function(x, ...) {
+  as.array(image_to_array(x, ...))
+}

man/layer_discretization.Rd

@@ -0,0 +1,12 @@
+test_that("DimExpr Ops keeps symbolic dims when R uses double scalars", {
+  skip_if_not(reticulate::py_module_available("jax"))
+
+  export <- reticulate::import("jax.export", convert = FALSE)
+  dim <- export$symbolic_shape("n")[[1]]
+
+  expr <- dim - 1 # 1 is a double in R; Ops method should coerce to int
+
+  expect_s3_class(expr, "jax._src.export.shape_poly._DimExpr")
+  expect_match(reticulate::py_str(expr), "n - 1")
+  expect_false(any(grepl("Array", class(expr))))
+})