diff --git a/neuralnetworks/aidl/Android.bp b/neuralnetworks/aidl/Android.bp
index b1860e2..ebf4654 100644
--- a/neuralnetworks/aidl/Android.bp
+++ b/neuralnetworks/aidl/Android.bp
@@ -16,6 +16,7 @@
     stability: "vintf",
     imports: [
         "android.hardware.common",
+        "android.hardware.graphics.common",
     ],
     backend: {
         java: {
diff --git a/neuralnetworks/aidl/aidl_api/android.hardware.neuralnetworks/current/android/hardware/neuralnetworks/Memory.aidl b/neuralnetworks/aidl/aidl_api/android.hardware.neuralnetworks/current/android/hardware/neuralnetworks/Memory.aidl
index 8207b25..37fa102 100644
--- a/neuralnetworks/aidl/aidl_api/android.hardware.neuralnetworks/current/android/hardware/neuralnetworks/Memory.aidl
+++ b/neuralnetworks/aidl/aidl_api/android.hardware.neuralnetworks/current/android/hardware/neuralnetworks/Memory.aidl
@@ -33,8 +33,8 @@
 
 package android.hardware.neuralnetworks;
 @VintfStability
-parcelable Memory {
-  android.hardware.common.NativeHandle handle;
-  long size;
-  String name;
+union Memory {
+  android.hardware.common.Ashmem ashmem;
+  android.hardware.common.MappableFile mappableFile;
+  android.hardware.graphics.common.HardwareBuffer hardwareBuffer;
 }
diff --git a/neuralnetworks/aidl/android/hardware/neuralnetworks/Memory.aidl b/neuralnetworks/aidl/android/hardware/neuralnetworks/Memory.aidl
index 870f0ae..244ac87 100644
--- a/neuralnetworks/aidl/android/hardware/neuralnetworks/Memory.aidl
+++ b/neuralnetworks/aidl/android/hardware/neuralnetworks/Memory.aidl
@@ -15,16 +15,26 @@
  */
 
 package android.hardware.neuralnetworks;
-import android.hardware.common.NativeHandle;
-import android.os.ParcelFileDescriptor;
+
+import android.hardware.common.Ashmem;
+import android.hardware.common.MappableFile;
+import android.hardware.graphics.common.HardwareBuffer;
 
 /**
- * A type that is used to pass pieces of shared memory between processes.
- * The type structure mimics hidl_memory type from HIDL.
+ * The different types of memory that can be shared across processes.
  */
 @VintfStability
-parcelable Memory {
-    NativeHandle handle;
-    long size;
-    String name;
+union Memory {
+    /**
+     * Ashmem hidl_memory type from HIDL.
+     */
+    Ashmem ashmem;
+    /**
+     * File that can be mapped.
+     */
+    MappableFile mappableFile;
+    /**
+     * AIDL representation of AHardwareBuffer.
+     */
+    HardwareBuffer hardwareBuffer;
 }