keras/src/backend/common/backend_utils.py/_vectorize_parse_input_dimensions
def _vectorize_parse_input_dimensions(
    args,
    input_core_dims,
):
    from keras.src import ops

    if len(args) != len(input_core_dims):
        raise TypeError(
            "wrong number of positional arguments: "
            f"expected {len(input_core_dims)}, got {len(args)}"
        )
    shapes = []
    dim_sizes: dict[str, int] = {}
    for arg, core_dims in zip(args, input_core_dims):
        _vectorize_update_dim_sizes(
            dim_sizes, arg.shape, core_dims, is_input=True
        )
        ndim = arg.ndim - len(core_dims)
        shapes.append(arg.shape[:ndim])
    broadcast_shape = shapes[0]
    for s in shapes:
        broadcast_shape = ops.broadcast_shapes(broadcast_shape, s)
    return broadcast_shape, dim_sizes