Dataset sharding warning

[andremfreitas]

Issue type

Support

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.16

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04.3 LTS

Mobile device

No response

Python version

3.10

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

A100 40GB

Current behavior?

Hi,

I am using a mirrored strategy for gpu parallelisation. This is how I prepare my dataset:

data = tf.data.Dataset.from_tensor_slices(gt_reshaped).batch(batch_size = batch_size).shuffle(buffer_size = n_ics)
dist_dataset = strategy.experimental_distribute_dataset(data)

Which is resulting in the warning below (see relevant log output).
This (closed) issue mentions something similar: #42146. I would expect this warning not to appear in tf2.16 anymore. I am not trying to do data sharding of a a file or files so this warning shouldn't be thrown in my opinion (and it's annoying). But please let me know if it makes sense to be thrown in a case like this.

Thanks,
Andre

Standalone code to reproduce the issue

Cannot produce a MWE at the moment.

Relevant log output

2024-05-02 17:21:33.678068: W tensorflow/core/grappler/optimizers/data/auto_shard.cc:766] AUTO sharding policy will apply DATA sharding policy as it failed to apply FILE sharding policy because of the following reason: Found an unshardable source dataset: name: "TensorSliceDataset/_1"
op: "TensorSliceDataset"
input: "Placeholder/_0"
attr {
  key: "Toutput_types"
  value {
    list {
      type: DT_COMPLEX128
    }
  }
}
attr {
  key: "_cardinality"
  value {
    i: 2560000
  }
}
attr {
  key: "is_files"
  value {
    b: false
  }
}
attr {
  key: "metadata"
  value {
    s: "\n\024TensorSliceDataset:0"
  }
}
attr {
  key: "output_shapes"
  value {
    list {
      shape {
        dim {
          size: 10
        }
        dim {
          size: 2
        }
      }
    }
  }
}

[sushreebarsa]

@andremfreitas Could you please provide a standalone code to replicate the issue reported here as it would help us to analyze the issue?
Thank you!

[github-actions]

This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you.