tf_agents.networks.utils.BatchSquash
Stay organized with collections
Save and categorize content based on your preferences.
Facilitates flattening and unflattening batch dims of a tensor.
tf_agents.networks.utils.BatchSquash(
batch_dims
)
Used in the notebooks
Exposes a pair of matched faltten and unflatten methods. After flattening only
1 batch dimension will be left. This facilitates evaluating networks that
expect inputs to have only 1 batch dimension.
Args |
batch_dims
|
Number of batch dimensions the flatten/unflatten ops should
handle.
|
Raises |
ValueError
|
if batch dims is negative.
|
Methods
flatten
View source
flatten(
tensor
)
Flattens and caches the tensor's batch_dims.
unflatten
View source
unflatten(
tensor
)
Unflattens the tensor's batch_dims using the cached shape.
Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates.
Last updated 2024-04-26 UTC.
[null,null,["Last updated 2024-04-26 UTC."],[],[],null,["# tf_agents.networks.utils.BatchSquash\n\n\u003cbr /\u003e\n\n|-----------------------------------------------------------------------------------------------------------------|\n| [View source on GitHub](https://github.com/tensorflow/agents/blob/v0.19.0/tf_agents/networks/utils.py#L49-L124) |\n\nFacilitates flattening and unflattening batch dims of a tensor. \n\n tf_agents.networks.utils.BatchSquash(\n batch_dims\n )\n\n### Used in the notebooks\n\n| Used in the tutorials |\n|-------------------------------------------------------------------------------|\n| - [Networks](https://www.tensorflow.org/agents/tutorials/8_networks_tutorial) |\n\nExposes a pair of matched faltten and unflatten methods. After flattening only\n1 batch dimension will be left. This facilitates evaluating networks that\nexpect inputs to have only 1 batch dimension.\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Args ---- ||\n|--------------|---------------------------------------------------------------------|\n| `batch_dims` | Number of batch dimensions the flatten/unflatten ops should handle. |\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Raises ------ ||\n|--------------|----------------------------|\n| `ValueError` | if batch dims is negative. |\n\n\u003cbr /\u003e\n\nMethods\n-------\n\n### `flatten`\n\n[View source](https://github.com/tensorflow/agents/blob/v0.19.0/tf_agents/networks/utils.py#L72-L106) \n\n flatten(\n tensor\n )\n\nFlattens and caches the tensor's batch_dims.\n\n### `unflatten`\n\n[View source](https://github.com/tensorflow/agents/blob/v0.19.0/tf_agents/networks/utils.py#L108-L124) \n\n unflatten(\n tensor\n )\n\nUnflattens the tensor's batch_dims using the cached shape."]]